diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java index 93fcc332a..1c8020ca5 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -42,6 +42,7 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.util.context.ContextView; /** * The Model Context Protocol (MCP) client implementation that provides asynchronous @@ -181,7 +182,7 @@ public class McpAsyncClient { * @param features the MCP Client supported features. responses against output * schemas. */ - McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout, + public McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout, JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features) { Assert.notNull(transport, "Transport must not be null"); @@ -317,13 +318,20 @@ public class McpAsyncClient { }; this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, transport.protocolVersions(), - initializationTimeout, ctx -> new McpClientSession(requestTimeout, transport, requestHandlers, - notificationHandlers, con -> con.contextWrite(ctx)), + initializationTimeout, + ctx -> buildClientSession(requestTimeout, transport, requestHandlers, notificationHandlers, ctx), postInitializationHook); this.transport.setExceptionHandler(this.initializer::handleException); } + protected McpClientSession buildClientSession(Duration requestTimeout, McpClientTransport transport, + Map> requestHandlers, Map notificationHandlers, + ContextView ctx) { + return new McpClientSession(requestTimeout, transport, requestHandlers, notificationHandlers, + con -> con.contextWrite(ctx)); + } + /** * Get the current initialization result. * @return the initialization result. diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java index 127d53337..7723af798 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java @@ -46,7 +46,7 @@ * @see McpSchema.Implementation * @see McpSchema.ClientCapabilities */ -class McpClientFeatures { +public class McpClientFeatures { /** * Asynchronous client features specification providing the capabilities and request @@ -64,7 +64,7 @@ class McpClientFeatures { * @param elicitationHandler the elicitation handler. * @param enableCallToolSchemaCaching whether to enable call tool schema caching. */ - record Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities, + public record Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities, Map roots, List, Mono>> toolsChangeConsumers, List, Mono>> resourcesChangeConsumers, List, Mono>> resourcesUpdateConsumers, diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java index 7fdaa8941..2f539f105 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java @@ -64,7 +64,7 @@ public class McpSyncClient implements AutoCloseable { // is not a requirement? private static final long DEFAULT_CLOSE_TIMEOUT_MS = 10_000L; - private final McpAsyncClient delegate; + protected final McpAsyncClient delegate; private final Supplier contextProvider; @@ -75,7 +75,7 @@ public class McpSyncClient implements AutoCloseable { * @param contextProvider the supplier of context before calling any non-blocking * operation on underlying delegate */ - McpSyncClient(McpAsyncClient delegate, Supplier contextProvider) { + public McpSyncClient(McpAsyncClient delegate, Supplier contextProvider) { Assert.notNull(delegate, "The delegate can not be null"); Assert.notNull(contextProvider, "The contextProvider can not be null"); this.delegate = delegate; diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java index 23285d514..ae6bc4395 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java @@ -91,35 +91,35 @@ public class McpAsyncServer { private static final Logger logger = LoggerFactory.getLogger(McpAsyncServer.class); - private final McpServerTransportProviderBase mcpTransportProvider; + protected final McpServerTransportProviderBase mcpTransportProvider; - private final McpJsonMapper jsonMapper; + protected final McpJsonMapper jsonMapper; - private final JsonSchemaValidator jsonSchemaValidator; + protected final JsonSchemaValidator jsonSchemaValidator; - private final McpSchema.ServerCapabilities serverCapabilities; + protected final McpSchema.ServerCapabilities serverCapabilities; - private final McpSchema.Implementation serverInfo; + protected final McpSchema.Implementation serverInfo; - private final String instructions; + protected final String instructions; - private final CopyOnWriteArrayList tools = new CopyOnWriteArrayList<>(); + protected final CopyOnWriteArrayList tools = new CopyOnWriteArrayList<>(); - private final ConcurrentHashMap resources = new ConcurrentHashMap<>(); + protected final ConcurrentHashMap resources = new ConcurrentHashMap<>(); - private final ConcurrentHashMap resourceTemplates = new ConcurrentHashMap<>(); + protected final ConcurrentHashMap resourceTemplates = new ConcurrentHashMap<>(); - private final ConcurrentHashMap prompts = new ConcurrentHashMap<>(); + protected final ConcurrentHashMap prompts = new ConcurrentHashMap<>(); // FIXME: this field is deprecated and should be remvoed together with the // broadcasting loggingNotification. private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG; - private final ConcurrentHashMap completions = new ConcurrentHashMap<>(); + protected final ConcurrentHashMap completions = new ConcurrentHashMap<>(); - private List protocolVersions; + protected List protocolVersions; - private McpUriTemplateManagerFactory uriTemplateManagerFactory = new DefaultMcpUriTemplateManagerFactory(); + protected McpUriTemplateManagerFactory uriTemplateManagerFactory = new DefaultMcpUriTemplateManagerFactory(); /** * Create a new McpAsyncServer with the given transport provider and capabilities. @@ -128,7 +128,7 @@ public class McpAsyncServer { * @param features The MCP server supported features. * @param jsonMapper The JsonMapper to use for JSON serialization/deserialization */ - McpAsyncServer(McpServerTransportProvider mcpTransportProvider, McpJsonMapper jsonMapper, + public McpAsyncServer(McpServerTransportProvider mcpTransportProvider, McpJsonMapper jsonMapper, McpServerFeatures.Async features, Duration requestTimeout, McpUriTemplateManagerFactory uriTemplateManagerFactory, JsonSchemaValidator jsonSchemaValidator) { this.mcpTransportProvider = mcpTransportProvider; @@ -153,7 +153,7 @@ public class McpAsyncServer { requestTimeout, transport, this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers)); } - McpAsyncServer(McpStreamableServerTransportProvider mcpTransportProvider, McpJsonMapper jsonMapper, + public McpAsyncServer(McpStreamableServerTransportProvider mcpTransportProvider, McpJsonMapper jsonMapper, McpServerFeatures.Async features, Duration requestTimeout, McpUriTemplateManagerFactory uriTemplateManagerFactory, JsonSchemaValidator jsonSchemaValidator) { this.mcpTransportProvider = mcpTransportProvider; @@ -178,7 +178,7 @@ public class McpAsyncServer { this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers)); } - private Map prepareNotificationHandlers(McpServerFeatures.Async features) { + protected Map prepareNotificationHandlers(McpServerFeatures.Async features) { Map notificationHandlers = new HashMap<>(); notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, (exchange, params) -> Mono.empty()); @@ -196,7 +196,7 @@ private Map prepareNotificationHandlers(McpServe return notificationHandlers; } - private Map> prepareRequestHandlers() { + protected Map> prepareRequestHandlers() { Map> requestHandlers = new HashMap<>(); // Initialize request handlers for standard MCP methods diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpServerFeatures.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpServerFeatures.java index fe0608b1c..fc888dbaf 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpServerFeatures.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpServerFeatures.java @@ -39,7 +39,7 @@ public class McpServerFeatures { * roots list changes * @param instructions The server instructions text */ - record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, Map prompts, @@ -59,7 +59,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s * the roots list changes * @param instructions The server instructions text */ - Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, Map prompts, @@ -101,7 +101,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s * @return a specification which is protected from blocking calls specified by the * user. */ - static Async fromSync(Sync syncSpec, boolean immediateExecution) { + public static Async fromSync(Sync syncSpec, boolean immediateExecution) { List tools = new ArrayList<>(); for (var tool : syncSpec.tools()) { tools.add(AsyncToolSpecification.fromSync(tool, immediateExecution)); @@ -153,7 +153,7 @@ static Async fromSync(Sync syncSpec, boolean immediateExecution) { * roots list changes * @param instructions The server instructions text */ - record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, @@ -173,7 +173,7 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se * the roots list changes * @param instructions The server instructions text */ - Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, @@ -354,7 +354,8 @@ public static Builder builder() { public record AsyncResourceSpecification(McpSchema.Resource resource, BiFunction> readHandler) { - static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, boolean immediateExecution) { + public static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, + boolean immediateExecution) { // FIXME: This is temporary, proper validation should be implemented if (resource == null) { return null; @@ -394,7 +395,7 @@ static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, b public record AsyncResourceTemplateSpecification(McpSchema.ResourceTemplate resourceTemplate, BiFunction> readHandler) { - static AsyncResourceTemplateSpecification fromSync(SyncResourceTemplateSpecification resource, + public static AsyncResourceTemplateSpecification fromSync(SyncResourceTemplateSpecification resource, boolean immediateExecution) { // FIXME: This is temporary, proper validation should be implemented if (resource == null) { @@ -442,7 +443,7 @@ static AsyncResourceTemplateSpecification fromSync(SyncResourceTemplateSpecifica public record AsyncPromptSpecification(McpSchema.Prompt prompt, BiFunction> promptHandler) { - static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt, boolean immediateExecution) { + public static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt, boolean immediateExecution) { // FIXME: This is temporary, proper validation should be implemented if (prompt == null) { return null; @@ -482,7 +483,7 @@ public record AsyncCompletionSpecification(McpSchema.CompleteReference reference * @return an asynchronous wrapper of the provided sync specification, or * {@code null} if input is null */ - static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion, + public static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion, boolean immediateExecution) { if (completion == null) { return null; diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java index c7a1fd0d7..59222caea 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java @@ -51,9 +51,9 @@ public class McpStatelessAsyncServer { private static final Logger logger = LoggerFactory.getLogger(McpStatelessAsyncServer.class); - private final McpStatelessServerTransport mcpTransportProvider; + protected final McpStatelessServerTransport mcpTransportProvider; - private final McpJsonMapper jsonMapper; + protected final McpJsonMapper jsonMapper; private final McpSchema.ServerCapabilities serverCapabilities; @@ -77,7 +77,7 @@ public class McpStatelessAsyncServer { private final JsonSchemaValidator jsonSchemaValidator; - McpStatelessAsyncServer(McpStatelessServerTransport mcpTransport, McpJsonMapper jsonMapper, + public McpStatelessAsyncServer(McpStatelessServerTransport mcpTransport, McpJsonMapper jsonMapper, McpStatelessServerFeatures.Async features, Duration requestTimeout, McpUriTemplateManagerFactory uriTemplateManagerFactory, JsonSchemaValidator jsonSchemaValidator) { this.mcpTransportProvider = mcpTransport; diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java index a15681ba5..46754fcdd 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java @@ -38,7 +38,7 @@ public class McpStatelessServerFeatures { * @param prompts The map of prompt specifications * @param instructions The server instructions text */ - record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, @@ -56,7 +56,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s * @param prompts The map of prompt specifications * @param instructions The server instructions text */ - Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, @@ -94,7 +94,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s * @return a specification which is protected from blocking calls specified by the * user. */ - static Async fromSync(Sync syncSpec, boolean immediateExecution) { + public static Async fromSync(Sync syncSpec, boolean immediateExecution) { List tools = new ArrayList<>(); for (var tool : syncSpec.tools()) { tools.add(AsyncToolSpecification.fromSync(tool, immediateExecution)); @@ -136,7 +136,7 @@ static Async fromSync(Sync syncSpec, boolean immediateExecution) { * @param prompts The map of prompt specifications * @param instructions The server instructions text */ - record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, @@ -154,7 +154,7 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se * @param prompts The map of prompt specifications * @param instructions The server instructions text */ - Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, + public Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities, List tools, Map resources, Map resourceTemplates, @@ -199,11 +199,11 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se public record AsyncToolSpecification(McpSchema.Tool tool, BiFunction> callHandler) { - static AsyncToolSpecification fromSync(SyncToolSpecification syncToolSpec) { + public static AsyncToolSpecification fromSync(SyncToolSpecification syncToolSpec) { return fromSync(syncToolSpec, false); } - static AsyncToolSpecification fromSync(SyncToolSpecification syncToolSpec, boolean immediate) { + public static AsyncToolSpecification fromSync(SyncToolSpecification syncToolSpec, boolean immediate) { // FIXME: This is temporary, proper validation should be implemented if (syncToolSpec == null) { @@ -291,7 +291,8 @@ public static Builder builder() { public record AsyncResourceSpecification(McpSchema.Resource resource, BiFunction> readHandler) { - static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, boolean immediateExecution) { + public static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, + boolean immediateExecution) { // FIXME: This is temporary, proper validation should be implemented if (resource == null) { return null; @@ -330,7 +331,7 @@ static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, b public record AsyncResourceTemplateSpecification(McpSchema.ResourceTemplate resourceTemplate, BiFunction> readHandler) { - static AsyncResourceTemplateSpecification fromSync(SyncResourceTemplateSpecification resource, + public static AsyncResourceTemplateSpecification fromSync(SyncResourceTemplateSpecification resource, boolean immediateExecution) { // FIXME: This is temporary, proper validation should be implemented if (resource == null) { @@ -362,7 +363,7 @@ static AsyncResourceTemplateSpecification fromSync(SyncResourceTemplateSpecifica public record AsyncPromptSpecification(McpSchema.Prompt prompt, BiFunction> promptHandler) { - static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt, boolean immediateExecution) { + public static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt, boolean immediateExecution) { // FIXME: This is temporary, proper validation should be implemented if (prompt == null) { return null; @@ -400,7 +401,7 @@ public record AsyncCompletionSpecification(McpSchema.CompleteReference reference * @return an asynchronous wrapper of the provided sync specification, or * {@code null} if input is null */ - static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion, + public static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion, boolean immediateExecution) { if (completion == null) { return null; diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessSyncServer.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessSyncServer.java index 6849eb8ed..376145255 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessSyncServer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpStatelessSyncServer.java @@ -27,7 +27,7 @@ public class McpStatelessSyncServer { private final boolean immediateExecution; - McpStatelessSyncServer(McpStatelessAsyncServer asyncServer, boolean immediateExecution) { + public McpStatelessSyncServer(McpStatelessAsyncServer asyncServer, boolean immediateExecution) { this.asyncServer = asyncServer; this.immediateExecution = immediateExecution; } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java index 10f0e5a31..c82b019b1 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java @@ -54,9 +54,9 @@ public class McpSyncServer { /** * The async server to wrap. */ - private final McpAsyncServer asyncServer; + protected final McpAsyncServer asyncServer; - private final boolean immediateExecution; + protected final boolean immediateExecution; /** * Creates a new synchronous server that wraps the provided async server.