Spring AI 源码解析：MCP链路调用流程及示例

public interface McpClient {


    static SyncSpec sync(McpClientTransport transport) {
       return new SyncSpec(transport);
    }


    static AsyncSpec async(McpClientTransport transport) {
       return new AsyncSpec(transport);
    }


    class SyncSpec {

       private final McpClientTransport transport;

       private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout

       private Duration initializationTimeout = Duration.ofSeconds(20);

       private ClientCapabilities capabilities;

       private Implementation clientInfo = new Implementation("Java SDK MCP Client", "1.0.0");

       private final Map<String, Root> roots = new HashMap<>();

       private final List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers = new ArrayList<>();

       private final List<Consumer<List<McpSchema.Resource>>> resourcesChangeConsumers = new ArrayList<>();

       private final List<Consumer<List<McpSchema.Prompt>>> promptsChangeConsumers = new ArrayList<>();

       private final List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers = new ArrayList<>();

       private Function<CreateMessageRequest, CreateMessageResult> samplingHandler;

       private SyncSpec(McpClientTransport transport) {
          Assert.notNull(transport, "Transport must not be null");
          this.transport = transport;
       }


       public SyncSpec requestTimeout(Duration requestTimeout) {
          Assert.notNull(requestTimeout, "Request timeout must not be null");
          this.requestTimeout = requestTimeout;
          return this;
       }


       public SyncSpec initializationTimeout(Duration initializationTimeout) {
          Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
          this.initializationTimeout = initializationTimeout;
          return this;
       }


       public SyncSpec capabilities(ClientCapabilities capabilities) {
          Assert.notNull(capabilities, "Capabilities must not be null");
          this.capabilities = capabilities;
          return this;
       }


       public SyncSpec clientInfo(Implementation clientInfo) {
          Assert.notNull(clientInfo, "Client info must not be null");
          this.clientInfo = clientInfo;
          return this;
       }


       public SyncSpec roots(List<Root> roots) {
          Assert.notNull(roots, "Roots must not be null");
          for (Root root : roots) {
             this.roots.put(root.uri(), root);
          }
          return this;
       }


       public SyncSpec roots(Root... roots) {
          Assert.notNull(roots, "Roots must not be null");
          for (Root root : roots) {
             this.roots.put(root.uri(), root);
          }
          return this;
       }


       public SyncSpec sampling(Function<CreateMessageRequest, CreateMessageResult> samplingHandler) {
          Assert.notNull(samplingHandler, "Sampling handler must not be null");
          this.samplingHandler = samplingHandler;
          return this;
       }


       public SyncSpec toolsChangeConsumer(Consumer<List<McpSchema.Tool>> toolsChangeConsumer) {
          Assert.notNull(toolsChangeConsumer, "Tools change consumer must not be null");
          this.toolsChangeConsumers.add(toolsChangeConsumer);
          return this;
       }


       public SyncSpec resourcesChangeConsumer(Consumer<List<McpSchema.Resource>> resourcesChangeConsumer) {
          Assert.notNull(resourcesChangeConsumer, "Resources change consumer must not be null");
          this.resourcesChangeConsumers.add(resourcesChangeConsumer);
          return this;
       }


       public SyncSpec promptsChangeConsumer(Consumer<List<McpSchema.Prompt>> promptsChangeConsumer) {
          Assert.notNull(promptsChangeConsumer, "Prompts change consumer must not be null");
          this.promptsChangeConsumers.add(promptsChangeConsumer);
          return this;
       }


       public SyncSpec loggingConsumer(Consumer<McpSchema.LoggingMessageNotification> loggingConsumer) {
          Assert.notNull(loggingConsumer, "Logging consumer must not be null");
          this.loggingConsumers.add(loggingConsumer);
          return this;
       }


       public SyncSpec loggingConsumers(List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers) {
          Assert.notNull(loggingConsumers, "Logging consumers must not be null");
          this.loggingConsumers.addAll(loggingConsumers);
          return this;
       }


       public McpSyncClient build() {
          McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities,
                this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.promptsChangeConsumers,
                this.loggingConsumers, this.samplingHandler);

          McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);

          return new McpSyncClient(
                new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures));
       }

    }


    class AsyncSpec {

       private final McpClientTransport transport;

       private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout

       private Duration initializationTimeout = Duration.ofSeconds(20);

       private ClientCapabilities capabilities;

       private Implementation clientInfo = new Implementation("Spring AI MCP Client", "0.3.1");

       private final Map<String, Root> roots = new HashMap<>();

       private final List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers = new ArrayList<>();

       private final List<Function<List<McpSchema.Resource>, Mono<Void>>> resourcesChangeConsumers = new ArrayList<>();

       private final List<Function<List<McpSchema.Prompt>, Mono<Void>>> promptsChangeConsumers = new ArrayList<>();

       private final List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers = new ArrayList<>();

       private Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler;

       private AsyncSpec(McpClientTransport transport) {
          Assert.notNull(transport, "Transport must not be null");
          this.transport = transport;
       }


       public AsyncSpec requestTimeout(Duration requestTimeout) {
          Assert.notNull(requestTimeout, "Request timeout must not be null");
          this.requestTimeout = requestTimeout;
          return this;
       }


       public AsyncSpec initializationTimeout(Duration initializationTimeout) {
          Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
          this.initializationTimeout = initializationTimeout;
          return this;
       }


       public AsyncSpec capabilities(ClientCapabilities capabilities) {
          Assert.notNull(capabilities, "Capabilities must not be null");
          this.capabilities = capabilities;
          return this;
       }


       public AsyncSpec clientInfo(Implementation clientInfo) {
          Assert.notNull(clientInfo, "Client info must not be null");
          this.clientInfo = clientInfo;
          return this;
       }


       public AsyncSpec roots(List<Root> roots) {
          Assert.notNull(roots, "Roots must not be null");
          for (Root root : roots) {
             this.roots.put(root.uri(), root);
          }
          return this;
       }


       public AsyncSpec roots(Root... roots) {
          Assert.notNull(roots, "Roots must not be null");
          for (Root root : roots) {
             this.roots.put(root.uri(), root);
          }
          return this;
       }


       public AsyncSpec sampling(Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler) {
          Assert.notNull(samplingHandler, "Sampling handler must not be null");
          this.samplingHandler = samplingHandler;
          return this;
       }


       public AsyncSpec toolsChangeConsumer(Function<List<McpSchema.Tool>, Mono<Void>> toolsChangeConsumer) {
          Assert.notNull(toolsChangeConsumer, "Tools change consumer must not be null");
          this.toolsChangeConsumers.add(toolsChangeConsumer);
          return this;
       }


       public AsyncSpec resourcesChangeConsumer(
             Function<List<McpSchema.Resource>, Mono<Void>> resourcesChangeConsumer) {
          Assert.notNull(resourcesChangeConsumer, "Resources change consumer must not be null");
          this.resourcesChangeConsumers.add(resourcesChangeConsumer);
          return this;
       }


       public AsyncSpec promptsChangeConsumer(Function<List<McpSchema.Prompt>, Mono<Void>> promptsChangeConsumer) {
          Assert.notNull(promptsChangeConsumer, "Prompts change consumer must not be null");
          this.promptsChangeConsumers.add(promptsChangeConsumer);
          return this;
       }


       public AsyncSpec loggingConsumer(Function<McpSchema.LoggingMessageNotification, Mono<Void>> loggingConsumer) {
          Assert.notNull(loggingConsumer, "Logging consumer must not be null");
          this.loggingConsumers.add(loggingConsumer);
          return this;
       }


       public AsyncSpec loggingConsumers(
             List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers) {
          Assert.notNull(loggingConsumers, "Logging consumers must not be null");
          this.loggingConsumers.addAll(loggingConsumers);
          return this;
       }


       public McpAsyncClient build() {
          return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout,
                new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
                      this.toolsChangeConsumers, this.resourcesChangeConsumers, this.promptsChangeConsumers,
                      this.loggingConsumers, this.samplingHandler));
       }

    }

}

public class McpAsyncClient {

    private static final Logger logger = LoggerFactory.getLogger(McpAsyncClient.class);

    private static TypeReference<Void> VOID_TYPE_REFERENCE = new TypeReference<>() {
    };

    protected final Sinks.One<McpSchema.InitializeResult> initializedSink = Sinks.one();

    private AtomicBoolean initialized = new AtomicBoolean(false);


    private final Duration initializationTimeout;


    private final McpClientSession mcpSession;


    private final McpSchema.ClientCapabilities clientCapabilities;


    private final McpSchema.Implementation clientInfo;


    private McpSchema.ServerCapabilities serverCapabilities;


    private McpSchema.Implementation serverInfo;


    private final ConcurrentHashMap<String, Root> roots;


    private Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler;


    private final McpTransport transport;


    private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);


    McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout,
          McpClientFeatures.Async features) {

       Assert.notNull(transport, "Transport must not be null");
       Assert.notNull(requestTimeout, "Request timeout must not be null");
       Assert.notNull(initializationTimeout, "Initialization timeout must not be null");

       this.clientInfo = features.clientInfo();
       this.clientCapabilities = features.clientCapabilities();
       this.transport = transport;
       this.roots = new ConcurrentHashMap<>(features.roots());
       this.initializationTimeout = initializationTimeout;

       // Request Handlers
       Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();

       // Roots List Request Handler
       if (this.clientCapabilities.roots() != null) {
          requestHandlers.put(McpSchema.METHOD_ROOTS_LIST, rootsListRequestHandler());
       }

       // Sampling Handler
       if (this.clientCapabilities.sampling() != null) {
          if (features.samplingHandler() == null) {
             throw new McpError("Sampling handler must not be null when client capabilities include sampling");
          }
          this.samplingHandler = features.samplingHandler();
          requestHandlers.put(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, samplingCreateMessageHandler());
       }

       // Notification Handlers
       Map<String, NotificationHandler> notificationHandlers = new HashMap<>();

       // Tools Change Notification
       List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumersFinal = new ArrayList<>();
       toolsChangeConsumersFinal
          .add((notification) -> Mono.fromRunnable(() -> logger.debug("Tools changed: {}", notification)));

       if (!Utils.isEmpty(features.toolsChangeConsumers())) {
          toolsChangeConsumersFinal.addAll(features.toolsChangeConsumers());
       }
       notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED,
             asyncToolsChangeNotificationHandler(toolsChangeConsumersFinal));

       // Resources Change Notification
       List<Function<List<McpSchema.Resource>, Mono<Void>>> resourcesChangeConsumersFinal = new ArrayList<>();
       resourcesChangeConsumersFinal
          .add((notification) -> Mono.fromRunnable(() -> logger.debug("Resources changed: {}", notification)));

       if (!Utils.isEmpty(features.resourcesChangeConsumers())) {
          resourcesChangeConsumersFinal.addAll(features.resourcesChangeConsumers());
       }

       notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED,
             asyncResourcesChangeNotificationHandler(resourcesChangeConsumersFinal));

       // Prompts Change Notification
       List<Function<List<McpSchema.Prompt>, Mono<Void>>> promptsChangeConsumersFinal = new ArrayList<>();
       promptsChangeConsumersFinal
          .add((notification) -> Mono.fromRunnable(() -> logger.debug("Prompts changed: {}", notification)));
       if (!Utils.isEmpty(features.promptsChangeConsumers())) {
          promptsChangeConsumersFinal.addAll(features.promptsChangeConsumers());
       }
       notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_PROMPTS_LIST_CHANGED,
             asyncPromptsChangeNotificationHandler(promptsChangeConsumersFinal));

       // Utility Logging Notification
       List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumersFinal = new ArrayList<>();
       loggingConsumersFinal.add((notification) -> Mono.fromRunnable(() -> logger.debug("Logging: {}", notification)));
       if (!Utils.isEmpty(features.loggingConsumers())) {
          loggingConsumersFinal.addAll(features.loggingConsumers());
       }
       notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_MESSAGE,
             asyncLoggingNotificationHandler(loggingConsumersFinal));

       this.mcpSession = new McpClientSession(requestTimeout, transport, requestHandlers, notificationHandlers);

    }


    public McpSchema.ServerCapabilities getServerCapabilities() {
       return this.serverCapabilities;
    }


    public McpSchema.Implementation getServerInfo() {
       return this.serverInfo;
    }


    public boolean isInitialized() {
       return this.initialized.get();
    }


    public ClientCapabilities getClientCapabilities() {
       return this.clientCapabilities;
    }


    public McpSchema.Implementation getClientInfo() {
       return this.clientInfo;
    }


    public void close() {
       this.mcpSession.close();
    }


    public Mono<Void> closeGracefully() {
       return this.mcpSession.closeGracefully();
    }

    // --------------------------
    // Initialization
    // --------------------------

    public Mono<McpSchema.InitializeResult> initialize() {

       String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);

       McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
             latestVersion,
             this.clientCapabilities,
             this.clientInfo); // @formatter:on

       Mono<McpSchema.InitializeResult> result = this.mcpSession.sendRequest(McpSchema.METHOD_INITIALIZE,
             initializeRequest, new TypeReference<McpSchema.InitializeResult>() {
             });

       return result.flatMap(initializeResult -> {

          this.serverCapabilities = initializeResult.capabilities();
          this.serverInfo = initializeResult.serverInfo();

          logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
                initializeResult.protocolVersion(), initializeResult.capabilities(), initializeResult.serverInfo(),
                initializeResult.instructions());

          if (!this.protocolVersions.contains(initializeResult.protocolVersion())) {
             return Mono.error(new McpError(
                   "Unsupported protocol version from the server: " + initializeResult.protocolVersion()));
          }

          return this.mcpSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null).doOnSuccess(v -> {
             this.initialized.set(true);
             this.initializedSink.tryEmitValue(initializeResult);
          }).thenReturn(initializeResult);
       });
    }


    private <T> Mono<T> withInitializationCheck(String actionName,
          Function<McpSchema.InitializeResult, Mono<T>> operation) {
       return this.initializedSink.asMono()
          .timeout(this.initializationTimeout)
          .onErrorResume(TimeoutException.class,
                ex -> Mono.error(new McpError("Client must be initialized before " + actionName)))
          .flatMap(operation);
    }

    // --------------------------
    // Basic Utilites
    // --------------------------


    public Mono<Object> ping() {
       return this.withInitializationCheck("pinging the server", initializedResult -> this.mcpSession
          .sendRequest(McpSchema.METHOD_PING, null, new TypeReference<Object>() {
          }));
    }

    // --------------------------
    // Roots
    // --------------------------

    public Mono<Void> addRoot(Root root) {

       if (root == null) {
          return Mono.error(new McpError("Root must not be null"));
       }

       if (this.clientCapabilities.roots() == null) {
          return Mono.error(new McpError("Client must be configured with roots capabilities"));
       }

       if (this.roots.containsKey(root.uri())) {
          return Mono.error(new McpError("Root with uri '" + root.uri() + "' already exists"));
       }

       this.roots.put(root.uri(), root);

       logger.debug("Added root: {}", root);

       if (this.clientCapabilities.roots().listChanged()) {
          if (this.isInitialized()) {
             return this.rootsListChangedNotification();
          }
          else {
             logger.warn("Client is not initialized, ignore sending a roots list changed notification");
          }
       }
       return Mono.empty();
    }


    public Mono<Void> removeRoot(String rootUri) {

       if (rootUri == null) {
          return Mono.error(new McpError("Root uri must not be null"));
       }

       if (this.clientCapabilities.roots() == null) {
          return Mono.error(new McpError("Client must be configured with roots capabilities"));
       }

       Root removed = this.roots.remove(rootUri);

       if (removed != null) {
          logger.debug("Removed Root: {}", rootUri);
          if (this.clientCapabilities.roots().listChanged()) {
             if (this.isInitialized()) {
                return this.rootsListChangedNotification();
             }
             else {
                logger.warn("Client is not initialized, ignore sending a roots list changed notification");
             }

          }
          return Mono.empty();
       }
       return Mono.error(new McpError("Root with uri '" + rootUri + "' not found"));
    }


    public Mono<Void> rootsListChangedNotification() {
       return this.withInitializationCheck("sending roots list changed notification",
             initResult -> this.mcpSession.sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED));
    }

    private RequestHandler<McpSchema.ListRootsResult> rootsListRequestHandler() {
       return params -> {
          @SuppressWarnings("unused")
          McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
                new TypeReference<McpSchema.PaginatedRequest>() {
                });

          List<Root> roots = this.roots.values().stream().toList();

          return Mono.just(new McpSchema.ListRootsResult(roots));
       };
    }

    // --------------------------
    // Sampling
    // --------------------------
    private RequestHandler<CreateMessageResult> samplingCreateMessageHandler() {
       return params -> {
          McpSchema.CreateMessageRequest request = transport.unmarshalFrom(params,
                new TypeReference<McpSchema.CreateMessageRequest>() {
                });

          return this.samplingHandler.apply(request);
       };
    }

    // --------------------------
    // Tools
    // --------------------------
    private static final TypeReference<McpSchema.CallToolResult> CALL_TOOL_RESULT_TYPE_REF = new TypeReference<>() {
    };

    private static final TypeReference<McpSchema.ListToolsResult> LIST_TOOLS_RESULT_TYPE_REF = new TypeReference<>() {
    };


    public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToolRequest) {
       return this.withInitializationCheck("calling tools", initializedResult -> {
          if (this.serverCapabilities.tools() == null) {
             return Mono.error(new McpError("Server does not provide tools capability"));
          }
          return this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
       });
    }


    public Mono<McpSchema.ListToolsResult> listTools() {
       return this.listTools(null);
    }


    public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
       return this.withInitializationCheck("listing tools", initializedResult -> {
          if (this.serverCapabilities.tools() == null) {
             return Mono.error(new McpError("Server does not provide tools capability"));
          }
          return this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
                LIST_TOOLS_RESULT_TYPE_REF);
       });
    }

    private NotificationHandler asyncToolsChangeNotificationHandler(
          List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers) {
       // TODO: params are not used yet
       return params -> this.listTools()
          .flatMap(listToolsResult -> Flux.fromIterable(toolsChangeConsumers)
             .flatMap(consumer -> consumer.apply(listToolsResult.tools()))
             .onErrorResume(error -> {
                logger.error("Error handling tools list change notification", error);
                return Mono.empty();
             })
             .then());
    }

    // --------------------------
    // Resources
    // --------------------------

    private static final TypeReference<McpSchema.ListResourcesResult> LIST_RESOURCES_RESULT_TYPE_REF = new TypeReference<>() {
    };

    private static final TypeReference<McpSchema.ReadResourceResult> READ_RESOURCE_RESULT_TYPE_REF = new TypeReference<>() {
    };

    private static final TypeReference<McpSchema.ListResourceTemplatesResult> LIST_RESOURCE_TEMPLATES_RESULT_TYPE_REF = new TypeReference<>() {
    };


    public Mono<McpSchema.ListResourcesResult> listResources() {
       return this.listResources(null);
    }


    public Mono<McpSchema.ListResourcesResult> listResources(String cursor) {
       return this.withInitializationCheck("listing resources", initializedResult -> {
          if (this.serverCapabilities.resources() == null) {
             return Mono.error(new McpError("Server does not provide the resources capability"));
          }
          return this.mcpSession.sendRequest(McpSchema.METHOD_RESOURCES_LIST, new McpSchema.PaginatedRequest(cursor),
                LIST_RESOURCES_RESULT_TYPE_REF);
       });
    }


    public Mono<McpSchema.ReadResourceResult> readResource(McpSchema.Resource resource) {
       return this.readResource(new McpSchema.ReadResourceRequest(resource.uri()));
    }


    public Mono<McpSchema.ReadResourceResult> readResource(McpSchema.ReadResourceRequest readResourceRequest) {
       return this.withInitializationCheck("reading resources", initializedResult -> {
          if (this.serverCapabilities.resources() == null) {
             return Mono.error(new McpError("Server does not provide the resources capability"));
          }
          return this.mcpSession.sendRequest(McpSchema.METHOD_RESOURCES_READ, readResourceRequest,
                READ_RESOURCE_RESULT_TYPE_REF);
       });
    }


    public Mono<McpSchema.ListResourceTemplatesResult> listResourceTemplates() {
       return this.listResourceTemplates(null);
    }


    public Mono<McpSchema.ListResourceTemplatesResult> listResourceTemplates(String cursor) {
       return this.withInitializationCheck("listing resource templates", initializedResult -> {
          if (this.serverCapabilities.resources() == null) {
             return Mono.error(new McpError("Server does not provide the resources capability"));
          }
          return this.mcpSession.sendRequest(McpSchema.METHOD_RESOURCES_TEMPLATES_LIST,
                new McpSchema.PaginatedRequest(cursor), LIST_RESOURCE_TEMPLATES_RESULT_TYPE_REF);
       });
    }


    public Mono<Void> subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
       return this.withInitializationCheck("subscribing to resources", initializedResult -> this.mcpSession
          .sendRequest(McpSchema.METHOD_RESOURCES_SUBSCRIBE, subscribeRequest, VOID_TYPE_REFERENCE));
    }


    public Mono<Void> unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
       return this.withInitializationCheck("unsubscribing from resources", initializedResult -> this.mcpSession
          .sendRequest(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, unsubscribeRequest, VOID_TYPE_REFERENCE));
    }

    private NotificationHandler asyncResourcesChangeNotificationHandler(
          List<Function<List<McpSchema.Resource>, Mono<Void>>> resourcesChangeConsumers) {
       return params -> listResources().flatMap(listResourcesResult -> Flux.fromIterable(resourcesChangeConsumers)
          .flatMap(consumer -> consumer.apply(listResourcesResult.resources()))
          .onErrorResume(error -> {
             logger.error("Error handling resources list change notification", error);
             return Mono.empty();
          })
          .then());
    }

    // --------------------------
    // Prompts
    // --------------------------
    private static final TypeReference<McpSchema.ListPromptsResult> LIST_PROMPTS_RESULT_TYPE_REF = new TypeReference<>() {
    };

    private static final TypeReference<McpSchema.GetPromptResult> GET_PROMPT_RESULT_TYPE_REF = new TypeReference<>() {
    };


    public Mono<ListPromptsResult> listPrompts() {
       return this.listPrompts(null);
    }


    public Mono<ListPromptsResult> listPrompts(String cursor) {
       return this.withInitializationCheck("listing prompts", initializedResult -> this.mcpSession
          .sendRequest(McpSchema.METHOD_PROMPT_LIST, new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF));
    }


    public Mono<GetPromptResult> getPrompt(GetPromptRequest getPromptRequest) {
       return this.withInitializationCheck("getting prompts", initializedResult -> this.mcpSession
          .sendRequest(McpSchema.METHOD_PROMPT_GET, getPromptRequest, GET_PROMPT_RESULT_TYPE_REF));
    }

    private NotificationHandler asyncPromptsChangeNotificationHandler(
          List<Function<List<McpSchema.Prompt>, Mono<Void>>> promptsChangeConsumers) {
       return params -> listPrompts().flatMap(listPromptsResult -> Flux.fromIterable(promptsChangeConsumers)
          .flatMap(consumer -> consumer.apply(listPromptsResult.prompts()))
          .onErrorResume(error -> {
             logger.error("Error handling prompts list change notification", error);
             return Mono.empty();
          })
          .then());
    }

    // --------------------------
    // Logging
    // --------------------------
    private NotificationHandler asyncLoggingNotificationHandler(
          List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumers) {

       return params -> {
          McpSchema.LoggingMessageNotification loggingMessageNotification = transport.unmarshalFrom(params,
                new TypeReference<McpSchema.LoggingMessageNotification>() {
                });

          return Flux.fromIterable(loggingConsumers)
             .flatMap(consumer -> consumer.apply(loggingMessageNotification))
             .then();
       };
    }


    public Mono<Void> setLoggingLevel(LoggingLevel loggingLevel) {
       if (loggingLevel == null) {
          return Mono.error(new McpError("Logging level must not be null"));
       }

       return this.withInitializationCheck("setting logging level", initializedResult -> {
          String levelName = this.transport.unmarshalFrom(loggingLevel, new TypeReference<String>() {
          });
          Map<String, Object> params = Map.of("level", levelName);
          return this.mcpSession.sendNotification(McpSchema.METHOD_LOGGING_SET_LEVEL, params);
       });
    }


    void setProtocolVersions(List<String> protocolVersions) {
       this.protocolVersions = protocolVersions;
    }

}

public class McpSyncClient implements AutoCloseable {

    private static final Logger logger = LoggerFactory.getLogger(McpSyncClient.class);

    // TODO: Consider providing a client config to set this properly
    // this is currently a concern only because AutoCloseable is used - perhaps it
    // is not a requirement?
    private static final long DEFAULT_CLOSE_TIMEOUT_MS = 10_000L;

    private final McpAsyncClient delegate;


    McpSyncClient(McpAsyncClient delegate) {
       Assert.notNull(delegate, "The delegate can not be null");
       this.delegate = delegate;
    }


    public McpSchema.ServerCapabilities getServerCapabilities() {
       return this.delegate.getServerCapabilities();
    }


    public McpSchema.Implementation getServerInfo() {
       return this.delegate.getServerInfo();
    }


    public ClientCapabilities getClientCapabilities() {
       return this.delegate.getClientCapabilities();
    }


    public McpSchema.Implementation getClientInfo() {
       return this.delegate.getClientInfo();
    }

    @Override
    public void close() {
       this.delegate.close();
    }

    public boolean closeGracefully() {
       try {
          this.delegate.closeGracefully().block(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS));
       }
       catch (RuntimeException e) {
          logger.warn("Client didn't close within timeout of {} ms.", DEFAULT_CLOSE_TIMEOUT_MS, e);
          return false;
       }
       return true;
    }


    public McpSchema.InitializeResult initialize() {
       // TODO: block takes no argument here as we assume the async client is
       // configured with a requestTimeout at all times
       return this.delegate.initialize().block();
    }


    public void rootsListChangedNotification() {
       this.delegate.rootsListChangedNotification().block();
    }


    public void addRoot(McpSchema.Root root) {
       this.delegate.addRoot(root).block();
    }


    public void removeRoot(String rootUri) {
       this.delegate.removeRoot(rootUri).block();
    }


    public Object ping() {
       return this.delegate.ping().block();
    }

    // --------------------------
    // Tools
    // --------------------------

    public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) {
       return this.delegate.callTool(callToolRequest).block();
    }


    public McpSchema.ListToolsResult listTools() {
       return this.delegate.listTools().block();
    }


    public McpSchema.ListToolsResult listTools(String cursor) {
       return this.delegate.listTools(cursor).block();
    }

    // --------------------------
    // Resources
    // --------------------------


    public McpSchema.ListResourcesResult listResources(String cursor) {
       return this.delegate.listResources(cursor).block();
    }


    public McpSchema.ListResourcesResult listResources() {
       return this.delegate.listResources().block();
    }


    public McpSchema.ReadResourceResult readResource(McpSchema.Resource resource) {
       return this.delegate.readResource(resource).block();
    }


    public McpSchema.ReadResourceResult readResource(McpSchema.ReadResourceRequest readResourceRequest) {
       return this.delegate.readResource(readResourceRequest).block();
    }


    public McpSchema.ListResourceTemplatesResult listResourceTemplates(String cursor) {
       return this.delegate.listResourceTemplates(cursor).block();
    }


    public McpSchema.ListResourceTemplatesResult listResourceTemplates() {
       return this.delegate.listResourceTemplates().block();
    }


    public void subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
       this.delegate.subscribeResource(subscribeRequest).block();
    }


    public void unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
       this.delegate.unsubscribeResource(unsubscribeRequest).block();
    }

    // --------------------------
    // Prompts
    // --------------------------
    public ListPromptsResult listPrompts(String cursor) {
       return this.delegate.listPrompts(cursor).block();
    }

    public ListPromptsResult listPrompts() {
       return this.delegate.listPrompts().block();
    }

    public GetPromptResult getPrompt(GetPromptRequest getPromptRequest) {
       return this.delegate.getPrompt(getPromptRequest).block();
    }


    public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) {
       this.delegate.setLoggingLevel(loggingLevel).block();
    }

}

public interface McpServer {


    static SyncSpecification sync(McpServerTransportProvider transportProvider) {
       return new SyncSpecification(transportProvider);
    }


    static AsyncSpecification async(McpServerTransportProvider transportProvider) {
       return new AsyncSpecification(transportProvider);
    }


    class AsyncSpecification {

       private static final McpSchema.Implementation DEFAULT_SERVER_INFO = new McpSchema.Implementation("mcp-server",
             "1.0.0");

       private final McpServerTransportProvider transportProvider;

       private ObjectMapper objectMapper;

       private McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO;

       private McpSchema.ServerCapabilities serverCapabilities;


       private final List<McpServerFeatures.AsyncToolSpecification> tools = new ArrayList<>();


       private final Map<String, McpServerFeatures.AsyncResourceSpecification> resources = new HashMap<>();

       private final List<ResourceTemplate> resourceTemplates = new ArrayList<>();


       private final Map<String, McpServerFeatures.AsyncPromptSpecification> prompts = new HashMap<>();

       private final List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeHandlers = new ArrayList<>();

       private AsyncSpecification(McpServerTransportProvider transportProvider) {
          Assert.notNull(transportProvider, "Transport provider must not be null");
          this.transportProvider = transportProvider;
       }


       public AsyncSpecification serverInfo(McpSchema.Implementation serverInfo) {
          Assert.notNull(serverInfo, "Server info must not be null");
          this.serverInfo = serverInfo;
          return this;
       }


       public AsyncSpecification serverInfo(String name, String version) {
          Assert.hasText(name, "Name must not be null or empty");
          Assert.hasText(version, "Version must not be null or empty");
          this.serverInfo = new McpSchema.Implementation(name, version);
          return this;
       }


       public AsyncSpecification capabilities(McpSchema.ServerCapabilities serverCapabilities) {
          Assert.notNull(serverCapabilities, "Server capabilities must not be null");
          this.serverCapabilities = serverCapabilities;
          return this;
       }


       public AsyncSpecification tool(McpSchema.Tool tool,
             BiFunction<McpAsyncServerExchange, Map<String, Object>, Mono<CallToolResult>> handler) {
          Assert.notNull(tool, "Tool must not be null");
          Assert.notNull(handler, "Handler must not be null");

          this.tools.add(new McpServerFeatures.AsyncToolSpecification(tool, handler));

          return this;
       }


       public AsyncSpecification tools(List<McpServerFeatures.AsyncToolSpecification> toolSpecifications) {
          Assert.notNull(toolSpecifications, "Tool handlers list must not be null");
          this.tools.addAll(toolSpecifications);
          return this;
       }


       public AsyncSpecification tools(McpServerFeatures.AsyncToolSpecification... toolSpecifications) {
          Assert.notNull(toolSpecifications, "Tool handlers list must not be null");
          for (McpServerFeatures.AsyncToolSpecification tool : toolSpecifications) {
             this.tools.add(tool);
          }
          return this;
       }


       public AsyncSpecification resources(
             Map<String, McpServerFeatures.AsyncResourceSpecification> resourceSpecifications) {
          Assert.notNull(resourceSpecifications, "Resource handlers map must not be null");
          this.resources.putAll(resourceSpecifications);
          return this;
       }


       public AsyncSpecification resources(List<McpServerFeatures.AsyncResourceSpecification> resourceSpecifications) {
          Assert.notNull(resourceSpecifications, "Resource handlers list must not be null");
          for (McpServerFeatures.AsyncResourceSpecification resource : resourceSpecifications) {
             this.resources.put(resource.resource().uri(), resource);
          }
          return this;
       }


       public AsyncSpecification resources(McpServerFeatures.AsyncResourceSpecification... resourceSpecifications) {
          Assert.notNull(resourceSpecifications, "Resource handlers list must not be null");
          for (McpServerFeatures.AsyncResourceSpecification resource : resourceSpecifications) {
             this.resources.put(resource.resource().uri(), resource);
          }
          return this;
       }


       public AsyncSpecification resourceTemplates(List<ResourceTemplate> resourceTemplates) {
          Assert.notNull(resourceTemplates, "Resource templates must not be null");
          this.resourceTemplates.addAll(resourceTemplates);
          return this;
       }


       public AsyncSpecification resourceTemplates(ResourceTemplate... resourceTemplates) {
          Assert.notNull(resourceTemplates, "Resource templates must not be null");
          for (ResourceTemplate resourceTemplate : resourceTemplates) {
             this.resourceTemplates.add(resourceTemplate);
          }
          return this;
       }


       public AsyncSpecification prompts(Map<String, McpServerFeatures.AsyncPromptSpecification> prompts) {
          Assert.notNull(prompts, "Prompts map must not be null");
          this.prompts.putAll(prompts);
          return this;
       }


       public AsyncSpecification prompts(List<McpServerFeatures.AsyncPromptSpecification> prompts) {
          Assert.notNull(prompts, "Prompts list must not be null");
          for (McpServerFeatures.AsyncPromptSpecification prompt : prompts) {
             this.prompts.put(prompt.prompt().name(), prompt);
          }
          return this;
       }


       public AsyncSpecification prompts(McpServerFeatures.AsyncPromptSpecification... prompts) {
          Assert.notNull(prompts, "Prompts list must not be null");
          for (McpServerFeatures.AsyncPromptSpecification prompt : prompts) {
             this.prompts.put(prompt.prompt().name(), prompt);
          }
          return this;
       }


       public AsyncSpecification rootsChangeHandler(
             BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>> handler) {
          Assert.notNull(handler, "Consumer must not be null");
          this.rootsChangeHandlers.add(handler);
          return this;
       }


       public AsyncSpecification rootsChangeHandlers(
             List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> handlers) {
          Assert.notNull(handlers, "Handlers list must not be null");
          this.rootsChangeHandlers.addAll(handlers);
          return this;
       }


       public AsyncSpecification rootsChangeHandlers(
             @SuppressWarnings("unchecked") BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>... handlers) {
          Assert.notNull(handlers, "Handlers list must not be null");
          return this.rootsChangeHandlers(Arrays.asList(handlers));
       }


       public AsyncSpecification objectMapper(ObjectMapper objectMapper) {
          Assert.notNull(objectMapper, "ObjectMapper must not be null");
          this.objectMapper = objectMapper;
          return this;
       }


       public McpAsyncServer build() {
          var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
                this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers);
          var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
          return new McpAsyncServer(this.transportProvider, mapper, features);
       }

    }


    class SyncSpecification {

       private static final McpSchema.Implementation DEFAULT_SERVER_INFO = new McpSchema.Implementation("mcp-server",
             "1.0.0");

       private final McpServerTransportProvider transportProvider;

       private ObjectMapper objectMapper;

       private McpSchema.Implementation serverInfo = DEFAULT_SERVER_INFO;

       private McpSchema.ServerCapabilities serverCapabilities;


       private final List<McpServerFeatures.SyncToolSpecification> tools = new ArrayList<>();


       private final Map<String, McpServerFeatures.SyncResourceSpecification> resources = new HashMap<>();

       private final List<ResourceTemplate> resourceTemplates = new ArrayList<>();


       private final Map<String, McpServerFeatures.SyncPromptSpecification> prompts = new HashMap<>();

       private final List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeHandlers = new ArrayList<>();

       private SyncSpecification(McpServerTransportProvider transportProvider) {
          Assert.notNull(transportProvider, "Transport provider must not be null");
          this.transportProvider = transportProvider;
       }


       public SyncSpecification serverInfo(McpSchema.Implementation serverInfo) {
          Assert.notNull(serverInfo, "Server info must not be null");
          this.serverInfo = serverInfo;
          return this;
       }


       public SyncSpecification serverInfo(String name, String version) {
          Assert.hasText(name, "Name must not be null or empty");
          Assert.hasText(version, "Version must not be null or empty");
          this.serverInfo = new McpSchema.Implementation(name, version);
          return this;
       }


       public SyncSpecification capabilities(McpSchema.ServerCapabilities serverCapabilities) {
          Assert.notNull(serverCapabilities, "Server capabilities must not be null");
          this.serverCapabilities = serverCapabilities;
          return this;
       }


       public SyncSpecification tool(McpSchema.Tool tool,
             BiFunction<McpSyncServerExchange, Map<String, Object>, McpSchema.CallToolResult> handler) {
          Assert.notNull(tool, "Tool must not be null");
          Assert.notNull(handler, "Handler must not be null");

          this.tools.add(new McpServerFeatures.SyncToolSpecification(tool, handler));

          return this;
       }


       public SyncSpecification tools(List<McpServerFeatures.SyncToolSpecification> toolSpecifications) {
          Assert.notNull(toolSpecifications, "Tool handlers list must not be null");
          this.tools.addAll(toolSpecifications);
          return this;
       }


       public SyncSpecification tools(McpServerFeatures.SyncToolSpecification... toolSpecifications) {
          Assert.notNull(toolSpecifications, "Tool handlers list must not be null");
          for (McpServerFeatures.SyncToolSpecification tool : toolSpecifications) {
             this.tools.add(tool);
          }
          return this;
       }


       public SyncSpecification resources(
             Map<String, McpServerFeatures.SyncResourceSpecification> resourceSpecifications) {
          Assert.notNull(resourceSpecifications, "Resource handlers map must not be null");
          this.resources.putAll(resourceSpecifications);
          return this;
       }


       public SyncSpecification resources(List<McpServerFeatures.SyncResourceSpecification> resourceSpecifications) {
          Assert.notNull(resourceSpecifications, "Resource handlers list must not be null");
          for (McpServerFeatures.SyncResourceSpecification resource : resourceSpecifications) {
             this.resources.put(resource.resource().uri(), resource);
          }
          return this;
       }


       public SyncSpecification resources(McpServerFeatures.SyncResourceSpecification... resourceSpecifications) {
          Assert.notNull(resourceSpecifications, "Resource handlers list must not be null");
          for (McpServerFeatures.SyncResourceSpecification resource : resourceSpecifications) {
             this.resources.put(resource.resource().uri(), resource);
          }
          return this;
       }


       public SyncSpecification resourceTemplates(List<ResourceTemplate> resourceTemplates) {
          Assert.notNull(resourceTemplates, "Resource templates must not be null");
          this.resourceTemplates.addAll(resourceTemplates);
          return this;
       }


       public SyncSpecification resourceTemplates(ResourceTemplate... resourceTemplates) {
          Assert.notNull(resourceTemplates, "Resource templates must not be null");
          for (ResourceTemplate resourceTemplate : resourceTemplates) {
             this.resourceTemplates.add(resourceTemplate);
          }
          return this;
       }


       public SyncSpecification prompts(Map<String, McpServerFeatures.SyncPromptSpecification> prompts) {
          Assert.notNull(prompts, "Prompts map must not be null");
          this.prompts.putAll(prompts);
          return this;
       }


       public SyncSpecification prompts(List<McpServerFeatures.SyncPromptSpecification> prompts) {
          Assert.notNull(prompts, "Prompts list must not be null");
          for (McpServerFeatures.SyncPromptSpecification prompt : prompts) {
             this.prompts.put(prompt.prompt().name(), prompt);
          }
          return this;
       }


       public SyncSpecification prompts(McpServerFeatures.SyncPromptSpecification... prompts) {
          Assert.notNull(prompts, "Prompts list must not be null");
          for (McpServerFeatures.SyncPromptSpecification prompt : prompts) {
             this.prompts.put(prompt.prompt().name(), prompt);
          }
          return this;
       }


       public SyncSpecification rootsChangeHandler(BiConsumer<McpSyncServerExchange, List<McpSchema.Root>> handler) {
          Assert.notNull(handler, "Consumer must not be null");
          this.rootsChangeHandlers.add(handler);
          return this;
       }


       public SyncSpecification rootsChangeHandlers(
             List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> handlers) {
          Assert.notNull(handlers, "Handlers list must not be null");
          this.rootsChangeHandlers.addAll(handlers);
          return this;
       }


       public SyncSpecification rootsChangeHandlers(
             BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>... handlers) {
          Assert.notNull(handlers, "Handlers list must not be null");
          return this.rootsChangeHandlers(List.of(handlers));
       }


       public SyncSpecification objectMapper(ObjectMapper objectMapper) {
          Assert.notNull(objectMapper, "ObjectMapper must not be null");
          this.objectMapper = objectMapper;
          return this;
       }


       public McpSyncServer build() {
          McpServerFeatures.Sync syncFeatures = new McpServerFeatures.Sync(this.serverInfo, this.serverCapabilities,
                this.tools, this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers);
          McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures);
          var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
          var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures);

          return new McpSyncServer(asyncServer);
       }

    }

}