From 8e11fa82ae0291d9b29566f5ab1299ddcb1ac639 Mon Sep 17 00:00:00 2001 From: Anurag Pant Date: Sun, 21 Sep 2025 20:12:32 -0700 Subject: [PATCH] feat(client): add client tool output schema validation and caching (#302) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add JSON schema validation for tool call results against output schemas - Implement automatic tool output schema caching during initialization - Add `enableCallToolSchemaCaching` configuration option to enable/disable schema caching - Add `JsonSchemaValidator` integration to McpClient builder APIs - Introduce post-initialization hook mechanism for performing operations after successful client initialization - Cache tool output schemas during `listTools` operations when caching is enabled - Validate tool results against cached schemas in `callTool` operations - Return error CallToolResult when validation fails - Add test coverage - Convert validateToolResult from Mono to synchronous method - Throw IllegalArgumentException on validation errors Signed-off-by: Christian Tzolov Co-authored-by: Anurag Pant Co-authored-by: Christian Tzolov Co-authored-by: Daniel Garnier-Moiroux Co-authored-by: Dariusz Jędrzejczyk --- .../client/LifecycleInitializer.java | 54 ++-- .../client/McpAsyncClient.java | 163 +++++++--- .../client/McpClient.java | 74 ++++- .../client/McpClientFeatures.java | 25 +- .../modelcontextprotocol/spec/McpSchema.java | 1 - ...pClientStreamableHttpAsyncClientTests.java | 5 +- ...nitializerPostInitializationHookTests.java | 280 ++++++++++++++++++ .../client/LifecycleInitializerTests.java | 70 ++--- .../client/McpAsyncClientTests.java | 156 +++++++++- .../client/McpClientProtocolVersionTests.java | 13 +- 10 files changed, 709 insertions(+), 132 deletions(-) create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerPostInitializationHookTests.java diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java index 2fc669c..f56c79a 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java @@ -11,14 +11,13 @@ import java.util.List; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import io.modelcontextprotocol.spec.McpClientSession; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException; import io.modelcontextprotocol.util.Assert; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import reactor.core.publisher.Sinks; import reactor.util.context.ContextView; @@ -99,21 +98,30 @@ class LifecycleInitializer { */ private final Duration initializationTimeout; + /** + * Post-initialization hook to perform additional operations after every successful + * initialization. + */ + private final Function> postInitializationHook; + public LifecycleInitializer(McpSchema.ClientCapabilities clientCapabilities, McpSchema.Implementation clientInfo, List protocolVersions, Duration initializationTimeout, - Function sessionSupplier) { + Function sessionSupplier, + Function> postInitializationHook) { Assert.notNull(sessionSupplier, "Session supplier must not be null"); Assert.notNull(clientCapabilities, "Client capabilities must not be null"); Assert.notNull(clientInfo, "Client info must not be null"); Assert.notEmpty(protocolVersions, "Protocol versions must not be empty"); Assert.notNull(initializationTimeout, "Initialization timeout must not be null"); + Assert.notNull(postInitializationHook, "Post-initialization hook must not be null"); this.sessionSupplier = sessionSupplier; this.clientCapabilities = clientCapabilities; this.clientInfo = clientInfo; this.protocolVersions = Collections.unmodifiableList(new ArrayList<>(protocolVersions)); this.initializationTimeout = initializationTimeout; + this.postInitializationHook = postInitializationHook; } /** @@ -148,10 +156,6 @@ class LifecycleInitializer { } - /** - * Default implementation of the {@link Initialization} interface that manages the MCP - * client initialization process. - */ private static class DefaultInitialization implements Initialization { /** @@ -199,29 +203,20 @@ class LifecycleInitializer { this.mcpClientSession.set(mcpClientSession); } - /** - * Returns a Mono that completes when the MCP client initialization is complete. - * This allows subscribers to wait for the initialization to finish before - * proceeding with further operations. - * @return A Mono that emits the result of the MCP initialization process - */ private Mono await() { return this.initSink.asMono(); } - /** - * Completes the initialization process with the given result. It caches the - * result and emits it to all subscribers waiting for the initialization to - * complete. - * @param initializeResult The result of the MCP initialization process - */ private void complete(McpSchema.InitializeResult initializeResult) { - // first ensure the result is cached - this.result.set(initializeResult); // inform all the subscribers waiting for the initialization this.initSink.emitValue(initializeResult, Sinks.EmitFailureHandler.FAIL_FAST); } + private void cacheResult(McpSchema.InitializeResult initializeResult) { + // first ensure the result is cached + this.result.set(initializeResult); + } + private void error(Throwable t) { this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST); } @@ -263,7 +258,7 @@ class LifecycleInitializer { } // Providing an empty operation since we are only interested in triggering // the implicit initialization step. - withIntitialization("re-initializing", result -> Mono.empty()).subscribe(); + this.withInitialization("re-initializing", result -> Mono.empty()).subscribe(); } } @@ -275,7 +270,7 @@ class LifecycleInitializer { * @param operation The operation to execute when the client is initialized * @return A Mono that completes with the result of the operation */ - public Mono withIntitialization(String actionName, Function> operation) { + public Mono withInitialization(String actionName, Function> operation) { return Mono.deferContextual(ctx -> { DefaultInitialization newInit = new DefaultInitialization(); DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit); @@ -283,8 +278,8 @@ class LifecycleInitializer { boolean needsToInitialize = previous == null; logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization"); - Mono initializationJob = needsToInitialize ? doInitialize(newInit, ctx) - : previous.await(); + Mono initializationJob = needsToInitialize + ? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await(); return initializationJob.map(initializeResult -> this.initializationRef.get()) .timeout(this.initializationTimeout) @@ -296,7 +291,9 @@ class LifecycleInitializer { }); } - private Mono doInitialize(DefaultInitialization initialization, ContextView ctx) { + private Mono doInitialize(DefaultInitialization initialization, + Function> postInitOperation, ContextView ctx) { + initialization.setMcpClientSession(this.sessionSupplier.apply(ctx)); McpClientSession mcpClientSession = initialization.mcpSession(); @@ -323,6 +320,9 @@ class LifecycleInitializer { return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null) .thenReturn(initializeResult); + }).flatMap(initializeResult -> { + initialization.cacheResult(initializeResult); + return postInitOperation.apply(initialization).thenReturn(initializeResult); }).doOnNext(initialization::complete).onErrorResume(ex -> { initialization.error(ex); return Mono.error(ex); 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 8d5bc34..53a05ae 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -15,12 +15,12 @@ import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - +import io.modelcontextprotocol.client.LifecycleInitializer.Initialization; import io.modelcontextprotocol.json.TypeRef; - +import io.modelcontextprotocol.json.schema.JsonSchemaValidator; import io.modelcontextprotocol.spec.McpClientSession; +import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler; +import io.modelcontextprotocol.spec.McpClientSession.RequestHandler; import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; @@ -35,10 +35,10 @@ import io.modelcontextprotocol.spec.McpSchema.LoggingLevel; import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification; import io.modelcontextprotocol.spec.McpSchema.PaginatedRequest; import io.modelcontextprotocol.spec.McpSchema.Root; -import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler; -import io.modelcontextprotocol.spec.McpClientSession.RequestHandler; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Utils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -75,6 +75,7 @@ import reactor.core.publisher.Mono; * @author Dariusz Jędrzejczyk * @author Christian Tzolov * @author Jihoon Kim + * @author Anurag Pant * @see McpClient * @see McpSchema * @see McpClientSession @@ -152,16 +153,33 @@ public class McpAsyncClient { */ private final LifecycleInitializer initializer; + /** + * JSON schema validator to use for validating tool responses against output schemas. + */ + private final JsonSchemaValidator jsonSchemaValidator; + + /** + * Cached tool output schemas. + */ + private final ConcurrentHashMap> toolsOutputSchemaCache; + + /** + * Whether to enable automatic schema caching during callTool operations. + */ + private final boolean enableCallToolSchemaCaching; + /** * Create a new McpAsyncClient with the given transport and session request-response * timeout. * @param transport the transport to use. * @param requestTimeout the session request-response timeout. * @param initializationTimeout the max timeout to await for the client-server - * @param features the MCP Client supported features. + * @param jsonSchemaValidator the JSON schema validator to use for validating tool + * @param features the MCP Client supported features. responses against output + * schemas. */ McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout, - McpClientFeatures.Async features) { + JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features) { Assert.notNull(transport, "Transport must not be null"); Assert.notNull(requestTimeout, "Request timeout must not be null"); @@ -171,6 +189,9 @@ public class McpAsyncClient { this.clientCapabilities = features.clientCapabilities(); this.transport = transport; this.roots = new ConcurrentHashMap<>(features.roots()); + this.jsonSchemaValidator = jsonSchemaValidator; + this.toolsOutputSchemaCache = new ConcurrentHashMap<>(); + this.enableCallToolSchemaCaching = features.enableCallToolSchemaCaching(); // Request Handlers Map> requestHandlers = new HashMap<>(); @@ -273,9 +294,30 @@ public class McpAsyncClient { notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_PROGRESS, asyncProgressNotificationHandler(progressConsumersFinal)); + Function> postInitializationHook = init -> { + + if (init.initializeResult().capabilities().tools() == null || !enableCallToolSchemaCaching) { + return Mono.empty(); + } + + return this.listToolsInternal(init, McpSchema.FIRST_PAGE).doOnNext(listToolsResult -> { + listToolsResult.tools() + .forEach(tool -> logger.debug("Tool {} schema: {}", tool.name(), tool.outputSchema())); + if (enableCallToolSchemaCaching && listToolsResult.tools() != null) { + // Cache tools output schema + listToolsResult.tools() + .stream() + .filter(tool -> tool.outputSchema() != null) + .forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema())); + } + }).then(); + }; + this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, transport.protocolVersions(), initializationTimeout, ctx -> new McpClientSession(requestTimeout, transport, requestHandlers, - notificationHandlers, con -> con.contextWrite(ctx))); + notificationHandlers, con -> con.contextWrite(ctx)), + postInitializationHook); + this.transport.setExceptionHandler(this.initializer::handleException); } @@ -387,7 +429,7 @@ public class McpAsyncClient { *

*/ public Mono initialize() { - return this.initializer.withIntitialization("by explicit API call", init -> Mono.just(init.initializeResult())); + return this.initializer.withInitialization("by explicit API call", init -> Mono.just(init.initializeResult())); } // -------------------------- @@ -399,7 +441,7 @@ public class McpAsyncClient { * @return A Mono that completes with the server's ping response */ public Mono ping() { - return this.initializer.withIntitialization("pinging the server", + return this.initializer.withInitialization("pinging the server", init -> init.mcpSession().sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF)); } @@ -480,7 +522,7 @@ public class McpAsyncClient { * @return A Mono that completes when the notification is sent. */ public Mono rootsListChangedNotification() { - return this.initializer.withIntitialization("sending roots list changed notification", + return this.initializer.withInitialization("sending roots list changed notification", init -> init.mcpSession().sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED)); } @@ -539,15 +581,45 @@ public class McpAsyncClient { * @see #listTools() */ public Mono callTool(McpSchema.CallToolRequest callToolRequest) { - return this.initializer.withIntitialization("calling tools", init -> { + return this.initializer.withInitialization("calling tool", init -> { if (init.initializeResult().capabilities().tools() == null) { return Mono.error(new IllegalStateException("Server does not provide tools capability")); } + return init.mcpSession() - .sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF); + .sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF) + .flatMap(result -> Mono.just(validateToolResult(callToolRequest.name(), result))); }); } + private McpSchema.CallToolResult validateToolResult(String toolName, McpSchema.CallToolResult result) { + + if (!this.enableCallToolSchemaCaching || result == null || result.isError() == Boolean.TRUE) { + // if tool schema caching is disabled or tool call resulted in an error - skip + // validation and return the result as it is + return result; + } + + Map optOutputSchema = this.toolsOutputSchemaCache.get(toolName); + + if (optOutputSchema == null) { + logger.warn( + "Calling a tool with no outputSchema is not expected to return result with structured content, but got: {}", + result.structuredContent()); + return result; + } + + // Validate the tool output against the cached output schema + var validation = this.jsonSchemaValidator.validate(optOutputSchema, result.structuredContent()); + + if (!validation.valid()) { + logger.warn("Tool call result validation failed: {}", validation.errorMessage()); + throw new IllegalArgumentException("Tool call result validation failed: " + validation.errorMessage()); + } + + return result; + } + /** * Retrieves the list of all tools provided by the server. * @return A Mono that emits the list of all tools result @@ -568,14 +640,26 @@ public class McpAsyncClient { * @return A Mono that emits the list of tools result */ public Mono listTools(String cursor) { - return this.initializer.withIntitialization("listing tools", init -> { - if (init.initializeResult().capabilities().tools() == null) { - return Mono.error(new IllegalStateException("Server does not provide tools capability")); - } - return init.mcpSession() - .sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor), - LIST_TOOLS_RESULT_TYPE_REF); - }); + return this.initializer.withInitialization("listing tools", init -> this.listToolsInternal(init, cursor)); + } + + private Mono listToolsInternal(Initialization init, String cursor) { + + if (init.initializeResult().capabilities().tools() == null) { + return Mono.error(new IllegalStateException("Server does not provide tools capability")); + } + return init.mcpSession() + .sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor), + LIST_TOOLS_RESULT_TYPE_REF) + .doOnNext(result -> { + if (this.enableCallToolSchemaCaching && result.tools() != null) { + // Cache tools output schema + result.tools() + .stream() + .filter(tool -> tool.outputSchema() != null) + .forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema())); + } + }); } private NotificationHandler asyncToolsChangeNotificationHandler( @@ -632,7 +716,7 @@ public class McpAsyncClient { * @see #readResource(McpSchema.Resource) */ public Mono listResources(String cursor) { - return this.initializer.withIntitialization("listing resources", init -> { + return this.initializer.withInitialization("listing resources", init -> { if (init.initializeResult().capabilities().resources() == null) { return Mono.error(new IllegalStateException("Server does not provide the resources capability")); } @@ -664,7 +748,7 @@ public class McpAsyncClient { * @see McpSchema.ReadResourceResult */ public Mono readResource(McpSchema.ReadResourceRequest readResourceRequest) { - return this.initializer.withIntitialization("reading resources", init -> { + return this.initializer.withInitialization("reading resources", init -> { if (init.initializeResult().capabilities().resources() == null) { return Mono.error(new IllegalStateException("Server does not provide the resources capability")); } @@ -702,7 +786,7 @@ public class McpAsyncClient { * @see McpSchema.ListResourceTemplatesResult */ public Mono listResourceTemplates(String cursor) { - return this.initializer.withIntitialization("listing resource templates", init -> { + return this.initializer.withInitialization("listing resource templates", init -> { if (init.initializeResult().capabilities().resources() == null) { return Mono.error(new IllegalStateException("Server does not provide the resources capability")); } @@ -722,7 +806,7 @@ public class McpAsyncClient { * @see #unsubscribeResource(McpSchema.UnsubscribeRequest) */ public Mono subscribeResource(McpSchema.SubscribeRequest subscribeRequest) { - return this.initializer.withIntitialization("subscribing to resources", init -> init.mcpSession() + return this.initializer.withInitialization("subscribing to resources", init -> init.mcpSession() .sendRequest(McpSchema.METHOD_RESOURCES_SUBSCRIBE, subscribeRequest, VOID_TYPE_REFERENCE)); } @@ -736,7 +820,7 @@ public class McpAsyncClient { * @see #subscribeResource(McpSchema.SubscribeRequest) */ public Mono unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) { - return this.initializer.withIntitialization("unsubscribing from resources", init -> init.mcpSession() + return this.initializer.withInitialization("unsubscribing from resources", init -> init.mcpSession() .sendRequest(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, unsubscribeRequest, VOID_TYPE_REFERENCE)); } @@ -802,7 +886,7 @@ public class McpAsyncClient { * @see #getPrompt(GetPromptRequest) */ public Mono listPrompts(String cursor) { - return this.initializer.withIntitialization("listing prompts", init -> init.mcpSession() + return this.initializer.withInitialization("listing prompts", init -> init.mcpSession() .sendRequest(McpSchema.METHOD_PROMPT_LIST, new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF)); } @@ -816,7 +900,7 @@ public class McpAsyncClient { * @see #listPrompts() */ public Mono getPrompt(GetPromptRequest getPromptRequest) { - return this.initializer.withIntitialization("getting prompts", init -> init.mcpSession() + return this.initializer.withInitialization("getting prompts", init -> init.mcpSession() .sendRequest(McpSchema.METHOD_PROMPT_GET, getPromptRequest, GET_PROMPT_RESULT_TYPE_REF)); } @@ -834,14 +918,6 @@ public class McpAsyncClient { // -------------------------- // Logging // -------------------------- - /** - * Create a notification handler for logging notifications from the server. This - * handler automatically distributes logging messages to all registered consumers. - * @param loggingConsumers List of consumers that will be notified when a logging - * message is received. Each consumer receives the logging message notification. - * @return A NotificationHandler that processes log notifications by distributing the - * message to all registered consumers - */ private NotificationHandler asyncLoggingNotificationHandler( List>> loggingConsumers) { @@ -867,7 +943,7 @@ public class McpAsyncClient { return Mono.error(new IllegalArgumentException("Logging level must not be null")); } - return this.initializer.withIntitialization("setting logging level", init -> { + return this.initializer.withInitialization("setting logging level", init -> { if (init.initializeResult().capabilities().logging() == null) { return Mono.error(new IllegalStateException("Server's Logging capabilities are not enabled!")); } @@ -876,15 +952,6 @@ public class McpAsyncClient { }); } - /** - * Create a notification handler for progress notifications from the server. This - * handler automatically distributes progress notifications to all registered - * consumers. - * @param progressConsumers List of consumers that will be notified when a progress - * message is received. Each consumer receives the progress notification. - * @return A NotificationHandler that processes progress notifications by distributing - * the message to all registered consumers - */ private NotificationHandler asyncProgressNotificationHandler( List>> progressConsumers) { @@ -924,7 +991,7 @@ public class McpAsyncClient { * @see McpSchema.CompleteResult */ public Mono completeCompletion(McpSchema.CompleteRequest completeRequest) { - return this.initializer.withIntitialization("complete completions", init -> init.mcpSession() + return this.initializer.withInitialization("complete completions", init -> init.mcpSession() .sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF)); } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java index 534879f..e39d43e 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java @@ -13,6 +13,7 @@ import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; +import io.modelcontextprotocol.json.schema.JsonSchemaValidator; import io.modelcontextprotocol.common.McpTransportContext; import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpSchema; @@ -99,6 +100,7 @@ import reactor.core.publisher.Mono; * * @author Christian Tzolov * @author Dariusz Jędrzejczyk + * @author Anurag Pant * @see McpAsyncClient * @see McpSyncClient * @see McpTransport @@ -187,6 +189,10 @@ public interface McpClient { private Supplier contextProvider = () -> McpTransportContext.EMPTY; + private JsonSchemaValidator jsonSchemaValidator; + + private boolean enableCallToolSchemaCaching = false; // Default to false + private SyncSpec(McpClientTransport transport) { Assert.notNull(transport, "Transport must not be null"); this.transport = transport; @@ -429,6 +435,32 @@ public interface McpClient { return this; } + /** + * Add a {@link JsonSchemaValidator} to validate the JSON structure of the + * structured output. + * @param jsonSchemaValidator A validator to validate the JSON structure of the + * structured output. Must not be null. + * @return This builder for method chaining + * @throws IllegalArgumentException if jsonSchemaValidator is null + */ + public SyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) { + Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null"); + this.jsonSchemaValidator = jsonSchemaValidator; + return this; + } + + /** + * Enables automatic schema caching during callTool operations. When a tool's + * output schema is not found in the cache, callTool will automatically fetch and + * cache all tool schemas via listTools. + * @param enableCallToolSchemaCaching true to enable, false to disable + * @return This builder instance for method chaining + */ + public SyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) { + this.enableCallToolSchemaCaching = enableCallToolSchemaCaching; + return this; + } + /** * Create an instance of {@link McpSyncClient} with the provided configurations or * sensible defaults. @@ -438,13 +470,13 @@ public interface McpClient { McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities, this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers, this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, this.samplingHandler, - this.elicitationHandler); + this.elicitationHandler, this.enableCallToolSchemaCaching); McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures); - return new McpSyncClient( - new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures), - this.contextProvider); + return new McpSyncClient(new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, + jsonSchemaValidator != null ? jsonSchemaValidator : JsonSchemaValidator.getDefault(), + asyncFeatures), this.contextProvider); } } @@ -495,6 +527,10 @@ public interface McpClient { private Function> elicitationHandler; + private JsonSchemaValidator jsonSchemaValidator; + + private boolean enableCallToolSchemaCaching = false; // Default to false + private AsyncSpec(McpClientTransport transport) { Assert.notNull(transport, "Transport must not be null"); this.transport = transport; @@ -741,17 +777,45 @@ public interface McpClient { return this; } + /** + * Sets the JSON schema validator to use for validating tool responses against + * output schemas. + * @param jsonSchemaValidator The validator to use. Must not be null. + * @return This builder instance for method chaining + * @throws IllegalArgumentException if jsonSchemaValidator is null + */ + public AsyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) { + Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null"); + this.jsonSchemaValidator = jsonSchemaValidator; + return this; + } + + /** + * Enables automatic schema caching during callTool operations. When a tool's + * output schema is not found in the cache, callTool will automatically fetch and + * cache all tool schemas via listTools. + * @param enableCallToolSchemaCaching true to enable, false to disable + * @return This builder instance for method chaining + */ + public AsyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) { + this.enableCallToolSchemaCaching = enableCallToolSchemaCaching; + return this; + } + /** * Create an instance of {@link McpAsyncClient} with the provided configurations * or sensible defaults. * @return a new instance of {@link McpAsyncClient}. */ public McpAsyncClient build() { + var jsonSchemaValidator = (this.jsonSchemaValidator != null) ? this.jsonSchemaValidator + : JsonSchemaValidator.getDefault(); return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout, + jsonSchemaValidator, new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers, this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, - this.samplingHandler, this.elicitationHandler)); + this.samplingHandler, this.elicitationHandler, this.enableCallToolSchemaCaching)); } } 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 3b65507..127d533 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpClientFeatures.java @@ -62,6 +62,7 @@ class McpClientFeatures { * @param progressConsumers the progress consumers. * @param samplingHandler the sampling handler. * @param elicitationHandler the elicitation handler. + * @param enableCallToolSchemaCaching whether to enable call tool schema caching. */ record Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities, Map roots, List, Mono>> toolsChangeConsumers, @@ -71,7 +72,8 @@ class McpClientFeatures { List>> loggingConsumers, List>> progressConsumers, Function> samplingHandler, - Function> elicitationHandler) { + Function> elicitationHandler, + boolean enableCallToolSchemaCaching) { /** * Create an instance and validate the arguments. @@ -84,6 +86,7 @@ class McpClientFeatures { * @param progressConsumers the progress consumers. * @param samplingHandler the sampling handler. * @param elicitationHandler the elicitation handler. + * @param enableCallToolSchemaCaching whether to enable call tool schema caching. */ public Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities, Map roots, @@ -94,7 +97,8 @@ class McpClientFeatures { List>> loggingConsumers, List>> progressConsumers, Function> samplingHandler, - Function> elicitationHandler) { + Function> elicitationHandler, + boolean enableCallToolSchemaCaching) { Assert.notNull(clientInfo, "Client info must not be null"); this.clientInfo = clientInfo; @@ -113,6 +117,7 @@ class McpClientFeatures { this.progressConsumers = progressConsumers != null ? progressConsumers : List.of(); this.samplingHandler = samplingHandler; this.elicitationHandler = elicitationHandler; + this.enableCallToolSchemaCaching = enableCallToolSchemaCaching; } /** @@ -129,7 +134,7 @@ class McpClientFeatures { Function> elicitationHandler) { this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers, resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler, - elicitationHandler); + elicitationHandler, false); } /** @@ -187,7 +192,8 @@ class McpClientFeatures { return new Async(syncSpec.clientInfo(), syncSpec.clientCapabilities(), syncSpec.roots(), toolsChangeConsumers, resourcesChangeConsumers, resourcesUpdateConsumers, promptsChangeConsumers, - loggingConsumers, progressConsumers, samplingHandler, elicitationHandler); + loggingConsumers, progressConsumers, samplingHandler, elicitationHandler, + syncSpec.enableCallToolSchemaCaching); } } @@ -205,6 +211,7 @@ class McpClientFeatures { * @param progressConsumers the progress consumers. * @param samplingHandler the sampling handler. * @param elicitationHandler the elicitation handler. + * @param enableCallToolSchemaCaching whether to enable call tool schema caching. */ public record Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities, Map roots, List>> toolsChangeConsumers, @@ -214,7 +221,8 @@ class McpClientFeatures { List> loggingConsumers, List> progressConsumers, Function samplingHandler, - Function elicitationHandler) { + Function elicitationHandler, + boolean enableCallToolSchemaCaching) { /** * Create an instance and validate the arguments. @@ -229,6 +237,7 @@ class McpClientFeatures { * @param progressConsumers the progress consumers. * @param samplingHandler the sampling handler. * @param elicitationHandler the elicitation handler. + * @param enableCallToolSchemaCaching whether to enable call tool schema caching. */ public Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities, Map roots, List>> toolsChangeConsumers, @@ -238,7 +247,8 @@ class McpClientFeatures { List> loggingConsumers, List> progressConsumers, Function samplingHandler, - Function elicitationHandler) { + Function elicitationHandler, + boolean enableCallToolSchemaCaching) { Assert.notNull(clientInfo, "Client info must not be null"); this.clientInfo = clientInfo; @@ -257,6 +267,7 @@ class McpClientFeatures { this.progressConsumers = progressConsumers != null ? progressConsumers : List.of(); this.samplingHandler = samplingHandler; this.elicitationHandler = elicitationHandler; + this.enableCallToolSchemaCaching = enableCallToolSchemaCaching; } /** @@ -272,7 +283,7 @@ class McpClientFeatures { Function elicitationHandler) { this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers, resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler, - elicitationHandler); + elicitationHandler, false); } } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpSchema.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpSchema.java index 8186bb1..12835a5 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpSchema.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpSchema.java @@ -17,7 +17,6 @@ import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonSubTypes; import com.fasterxml.jackson.annotation.JsonTypeInfo; -import com.fasterxml.jackson.annotation.JsonTypeInfo.As; import io.modelcontextprotocol.json.McpJsonMapper; import io.modelcontextprotocol.json.TypeRef; import io.modelcontextprotocol.util.Assert; diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java index 1e015ca..c4157bc 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java @@ -4,15 +4,14 @@ package io.modelcontextprotocol.client; +import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; +import io.modelcontextprotocol.spec.McpClientTransport; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Timeout; import org.testcontainers.containers.GenericContainer; import org.testcontainers.containers.wait.strategy.Wait; -import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; -import io.modelcontextprotocol.spec.McpClientTransport; - @Timeout(15) public class HttpClientStreamableHttpAsyncClientTests extends AbstractMcpAsyncClientTests { diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerPostInitializationHookTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerPostInitializationHookTests.java new file mode 100644 index 0000000..6f7390f --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerPostInitializationHookTests.java @@ -0,0 +1,280 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + +package io.modelcontextprotocol.client; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +import io.modelcontextprotocol.client.LifecycleInitializer.Initialization; +import io.modelcontextprotocol.spec.McpClientSession; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; +import reactor.test.StepVerifier; +import reactor.util.context.ContextView; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link LifecycleInitializer} postInitializationHook functionality. + * + * @author Christian Tzolov + */ +class LifecycleInitializerPostInitializationHookTests { + + private static final Duration INITIALIZATION_TIMEOUT = Duration.ofSeconds(5); + + private static final McpSchema.ClientCapabilities CLIENT_CAPABILITIES = McpSchema.ClientCapabilities.builder() + .build(); + + private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0"); + + private static final List PROTOCOL_VERSIONS = List.of("1.0.0", "2.0.0"); + + private static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult("2.0.0", + McpSchema.ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), + "Test instructions"); + + @Mock + private McpClientSession mockClientSession; + + @Mock + private Function mockSessionSupplier; + + @Mock + private Function> mockPostInitializationHook; + + private LifecycleInitializer initializer; + + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + + when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty()); + when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession); + when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())) + .thenReturn(Mono.just(MOCK_INIT_RESULT)); + when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any())) + .thenReturn(Mono.empty()); + when(mockClientSession.closeGracefully()).thenReturn(Mono.empty()); + + initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, + INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook); + } + + @Test + void shouldInvokePostInitializationHook() { + AtomicReference capturedInit = new AtomicReference<>(); + + when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> { + capturedInit.set(invocation.getArgument(0)); + return Mono.empty(); + }); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectNext(MOCK_INIT_RESULT) + .verifyComplete(); + + // Verify hook was called + verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class)); + + // Verify the hook received correct initialization data + assertThat(capturedInit.get()).isNotNull(); + assertThat(capturedInit.get().mcpSession()).isEqualTo(mockClientSession); + assertThat(capturedInit.get().initializeResult()).isEqualTo(MOCK_INIT_RESULT); + } + + @Test + void shouldInvokePostInitializationHookOnlyOnce() { + // First initialization + StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1"))) + .expectNext("result1") + .verifyComplete(); + + // Second call should reuse initialization and NOT call hook again + StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2"))) + .expectNext("result2") + .verifyComplete(); + + // Hook should only be called once + verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class)); + } + + @Test + void shouldInvokePostInitializationHookOnlyOnceWithConcurrentRequests() { + AtomicInteger hookInvocationCount = new AtomicInteger(0); + + when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> { + hookInvocationCount.incrementAndGet(); + return Mono.empty(); + }); + + // Start multiple concurrent initializations + Mono init1 = initializer.withInitialization("test1", init -> Mono.just("result1")) + .subscribeOn(Schedulers.parallel()); + Mono init2 = initializer.withInitialization("test2", init -> Mono.just("result2")) + .subscribeOn(Schedulers.parallel()); + Mono init3 = initializer.withInitialization("test3", init -> Mono.just("result3")) + .subscribeOn(Schedulers.parallel()); + + // TODO: can we assume the order of results? + StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> { + assertThat(tuple.getT1()).isEqualTo("result1"); + assertThat(tuple.getT2()).isEqualTo("result2"); + assertThat(tuple.getT3()).isEqualTo("result3"); + }).verifyComplete(); + + // Hook should only be called once despite concurrent requests + assertThat(hookInvocationCount.get()).isEqualTo(1); + } + + @Test + void shouldFailInitializationWhenPostInitializationHookFails() { + RuntimeException hookError = new RuntimeException("Post-initialization hook failed"); + when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.error(hookError)); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectErrorMatches(ex -> ex instanceof RuntimeException && ex.getCause() == hookError) + .verify(); + + // Verify initialization was not completed + assertThat(initializer.isInitialized()).isFalse(); + assertThat(initializer.currentInitializationResult()).isNull(); + + // Verify the hook was called + verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class)); + } + + @Test + void shouldNotInvokePostInitializationHookWhenInitializationFails() { + when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())) + .thenReturn(Mono.error(new RuntimeException("Initialization failed"))); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectError(RuntimeException.class) + .verify(); + + // Hook should NOT be called when initialization fails + verify(mockPostInitializationHook, never()).apply(any(Initialization.class)); + } + + @Test + void shouldNotInvokePostInitializationHookWhenNotificationFails() { + when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any())) + .thenReturn(Mono.error(new RuntimeException("Notification failed"))); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectError(RuntimeException.class) + .verify(); + + // Hook should NOT be called when notification fails + verify(mockPostInitializationHook, never()).apply(any(Initialization.class)); + } + + @Test + void shouldInvokePostInitializationHookAgainAfterReinitialization() { + AtomicInteger hookInvocationCount = new AtomicInteger(0); + + when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> { + hookInvocationCount.incrementAndGet(); + return Mono.empty(); + }); + + // First initialization + StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1"))) + .expectNext("result1") + .verifyComplete(); + + assertThat(hookInvocationCount.get()).isEqualTo(1); + + // Simulate transport session exception to trigger re-initialization + initializer.handleException(new McpTransportSessionNotFoundException("Session lost")); + + // Hook should be called twice (once for each initialization) + assertThat(hookInvocationCount.get()).isEqualTo(2); + } + + @Test + void shouldAllowPostInitializationHookToPerformAsyncOperations() { + AtomicInteger operationCount = new AtomicInteger(0); + + when(mockPostInitializationHook.apply(any(Initialization.class))) + .thenReturn(Mono.fromRunnable(() -> operationCount.incrementAndGet()).then()); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectNext(MOCK_INIT_RESULT) + .verifyComplete(); + + // Verify the async operation was executed + assertThat(operationCount.get()).isEqualTo(1); + verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class)); + } + + @Test + void shouldProvideCorrectInitializationDataToHook() { + AtomicReference capturedSession = new AtomicReference<>(); + AtomicReference capturedResult = new AtomicReference<>(); + + when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> { + Initialization init = invocation.getArgument(0); + capturedSession.set(init.mcpSession()); + capturedResult.set(init.initializeResult()); + return Mono.empty(); + }); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectNext(MOCK_INIT_RESULT) + .verifyComplete(); + + // Verify the hook received the correct session and result + assertThat(capturedSession.get()).isEqualTo(mockClientSession); + assertThat(capturedResult.get()).isEqualTo(MOCK_INIT_RESULT); + assertThat(capturedResult.get().protocolVersion()).isEqualTo("2.0.0"); + assertThat(capturedResult.get().serverInfo().name()).isEqualTo("test-server"); + } + + @Test + void shouldInvokePostInitializationHookAfterSuccessfulInitialization() { + AtomicReference notificationSent = new AtomicReference<>(false); + AtomicReference hookCalledAfterNotification = new AtomicReference<>(false); + + when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any())) + .thenAnswer(invocation -> { + notificationSent.set(true); + return Mono.empty(); + }); + + when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> { + // Due to flatMap chaining in doInitialize, if the hook is called, + // the notification must have been sent first + hookCalledAfterNotification.set(notificationSent.get()); + return Mono.empty(); + }); + + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) + .expectNext(MOCK_INIT_RESULT) + .verifyComplete(); + + // Verify the hook was called and notification was already sent at that point + assertThat(hookCalledAfterNotification.get()).isTrue(); + verify(mockClientSession).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()); + verify(mockPostInitializationHook).apply(any(Initialization.class)); + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java index 19de14c..787ee94 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java @@ -10,14 +10,14 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; +import io.modelcontextprotocol.client.LifecycleInitializer.Initialization; +import io.modelcontextprotocol.spec.McpClientSession; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; - -import io.modelcontextprotocol.spec.McpClientSession; -import io.modelcontextprotocol.spec.McpSchema; -import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; import reactor.test.StepVerifier; @@ -58,12 +58,16 @@ class LifecycleInitializerTests { @Mock private Function mockSessionSupplier; + @Mock + private Function> mockPostInitializationHook; + private LifecycleInitializer initializer; @BeforeEach void setUp() { MockitoAnnotations.openMocks(this); + when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty()); when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession); when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())) .thenReturn(Mono.just(MOCK_INIT_RESULT)); @@ -72,45 +76,45 @@ class LifecycleInitializerTests { when(mockClientSession.closeGracefully()).thenReturn(Mono.empty()); initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, - INITIALIZATION_TIMEOUT, mockSessionSupplier); + INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook); } @Test void constructorShouldValidateParameters() { assertThatThrownBy(() -> new LifecycleInitializer(null, CLIENT_INFO, PROTOCOL_VERSIONS, INITIALIZATION_TIMEOUT, - mockSessionSupplier)) + mockSessionSupplier, mockPostInitializationHook)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("Client capabilities must not be null"); assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, null, PROTOCOL_VERSIONS, - INITIALIZATION_TIMEOUT, mockSessionSupplier)) + INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("Client info must not be null"); assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, null, - INITIALIZATION_TIMEOUT, mockSessionSupplier)) + INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("Protocol versions must not be empty"); assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, List.of(), - INITIALIZATION_TIMEOUT, mockSessionSupplier)) + INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("Protocol versions must not be empty"); assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, null, - mockSessionSupplier)) + mockSessionSupplier, mockPostInitializationHook)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("Initialization timeout must not be null"); assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, - INITIALIZATION_TIMEOUT, null)) + INITIALIZATION_TIMEOUT, null, mockPostInitializationHook)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("Session supplier must not be null"); } @Test void shouldInitializeSuccessfully() { - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .assertNext(result -> { assertThat(result).isEqualTo(MOCK_INIT_RESULT); assertThat(initializer.isInitialized()).isTrue(); @@ -132,7 +136,7 @@ class LifecycleInitializerTests { return Mono.just(MOCK_INIT_RESULT); }); - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .assertNext(result -> { assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0"); // Latest // version @@ -152,7 +156,7 @@ class LifecycleInitializerTests { when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())) .thenReturn(Mono.just(unsupportedResult)); - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectError(RuntimeException.class) .verify(); @@ -167,13 +171,13 @@ class LifecycleInitializerTests { Duration SLOW_RESPONSE_DELAY = Duration.ofSeconds(5); LifecycleInitializer shortTimeoutInitializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, - PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier); + PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier, mockPostInitializationHook); when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())) .thenReturn(Mono.just(MOCK_INIT_RESULT).delayElement(SLOW_RESPONSE_DELAY, virtualTimeScheduler)); StepVerifier - .withVirtualTime(() -> shortTimeoutInitializer.withIntitialization("test", + .withVirtualTime(() -> shortTimeoutInitializer.withInitialization("test", init -> Mono.just(init.initializeResult())), () -> virtualTimeScheduler, Long.MAX_VALUE) .expectSubscription() .expectNoEvent(INITIALIZE_TIMEOUT) @@ -184,12 +188,12 @@ class LifecycleInitializerTests { @Test void shouldReuseExistingInitialization() { // First initialization - StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1"))) + StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1"))) .expectNext("result1") .verifyComplete(); // Second call should reuse the same initialization - StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2"))) + StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2"))) .expectNext("result2") .verifyComplete(); @@ -209,11 +213,11 @@ class LifecycleInitializerTests { // Start multiple concurrent initializations using subscribeOn with parallel // scheduler - Mono init1 = initializer.withIntitialization("test1", init -> Mono.just("result1")) + Mono init1 = initializer.withInitialization("test1", init -> Mono.just("result1")) .subscribeOn(Schedulers.parallel()); - Mono init2 = initializer.withIntitialization("test2", init -> Mono.just("result2")) + Mono init2 = initializer.withInitialization("test2", init -> Mono.just("result2")) .subscribeOn(Schedulers.parallel()); - Mono init3 = initializer.withIntitialization("test3", init -> Mono.just("result3")) + Mono init3 = initializer.withInitialization("test3", init -> Mono.just("result3")) .subscribeOn(Schedulers.parallel()); StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> { @@ -235,7 +239,7 @@ class LifecycleInitializerTests { // succeeds on the second call .thenReturn(Mono.just(MOCK_INIT_RESULT)); - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectError(RuntimeException.class) .verify(); @@ -244,7 +248,7 @@ class LifecycleInitializerTests { // The initializer can recover from previous errors StepVerifier - .create(initializer.withIntitialization("successful init", init -> Mono.just(init.initializeResult()))) + .create(initializer.withInitialization("successful init", init -> Mono.just(init.initializeResult()))) .expectNext(MOCK_INIT_RESULT) .verifyComplete(); @@ -255,7 +259,7 @@ class LifecycleInitializerTests { @Test void shouldHandleTransportSessionNotFoundException() { // successful initialization first - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectNext(MOCK_INIT_RESULT) .verifyComplete(); @@ -277,7 +281,7 @@ class LifecycleInitializerTests { @Test void shouldHandleOtherExceptions() { // Simulate a successful initialization first - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectNext(MOCK_INIT_RESULT) .verifyComplete(); @@ -295,7 +299,7 @@ class LifecycleInitializerTests { @Test void shouldCloseGracefully() { - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectNext(MOCK_INIT_RESULT) .verifyComplete(); @@ -307,7 +311,7 @@ class LifecycleInitializerTests { @Test void shouldCloseImmediately() { - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectNext(MOCK_INIT_RESULT) .verifyComplete(); @@ -342,7 +346,7 @@ class LifecycleInitializerTests { new McpSchema.Implementation("test-server", "1.0.0"), "Test instructions")); }); - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .assertNext(result -> { // Latest from new versions assertThat(capturedRequest.get().protocolVersion()).isEqualTo("4.0.0"); @@ -363,7 +367,7 @@ class LifecycleInitializerTests { }); StepVerifier - .create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())) + .create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())) .contextWrite(Context.of(contextKey, contextValue))) .expectNext(MOCK_INIT_RESULT) .verifyComplete(); @@ -374,7 +378,7 @@ class LifecycleInitializerTests { @Test void shouldProvideAccessToMcpSessionAndInitializeResult() { - StepVerifier.create(initializer.withIntitialization("test", init -> { + StepVerifier.create(initializer.withInitialization("test", init -> { assertThat(init.mcpSession()).isEqualTo(mockClientSession); assertThat(init.initializeResult()).isEqualTo(MOCK_INIT_RESULT); return Mono.just("success"); @@ -386,7 +390,7 @@ class LifecycleInitializerTests { when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any())) .thenReturn(Mono.error(new RuntimeException("Notification failed"))); - StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) + StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))) .expectError(RuntimeException.class) .verify(); @@ -403,7 +407,7 @@ class LifecycleInitializerTests { @Test void shouldReinitializeAfterTransportSessionException() { // First initialization - StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1"))) + StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1"))) .expectNext("result1") .verifyComplete(); @@ -411,7 +415,7 @@ class LifecycleInitializerTests { initializer.handleException(new McpTransportSessionNotFoundException("Session lost")); // Should be able to initialize again - StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2"))) + StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2"))) .expectNext("result2") .verifyComplete(); diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java index 3e29e89..970d8f2 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java @@ -10,12 +10,20 @@ import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.ProtocolVersions; import org.junit.jupiter.api.Test; -import reactor.core.publisher.Mono; +import com.fasterxml.jackson.core.JsonProcessingException; + +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import java.util.List; +import java.util.Map; +import java.util.Objects; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; class McpAsyncClientTests { @@ -24,6 +32,7 @@ class McpAsyncClientTests { "1.0.0"); public static final McpSchema.ServerCapabilities MOCK_SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder() + .tools(true) .build(); public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult( @@ -31,6 +40,91 @@ class McpAsyncClientTests { private static final String CONTEXT_KEY = "context.key"; + private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput) + throws JsonProcessingException { + + // Create tool with or without output schema + Map inputSchemaMap = Map.of("type", "object", "properties", + Map.of("expression", Map.of("type", "string")), "required", List.of("expression")); + + McpSchema.JsonSchema inputSchema = new McpSchema.JsonSchema("object", inputSchemaMap, null, null, null, null); + McpSchema.Tool.Builder toolBuilder = McpSchema.Tool.builder() + .name("calculator") + .description("Performs mathematical calculations") + .inputSchema(inputSchema); + + if (hasOutputSchema) { + Map outputSchema = Map.of("type", "object", "properties", + Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required", + List.of("result", "operation")); + toolBuilder.outputSchema(outputSchema); + } + + McpSchema.Tool calculatorTool = toolBuilder.build(); + McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(calculatorTool), null); + + // Create call tool result - valid or invalid based on parameter + Map structuredContent = invalidOutput ? Map.of("result", "5", "operation", "add") + : Map.of("result", 5, "operation", "add"); + + McpSchema.CallToolResult mockCallToolResult = McpSchema.CallToolResult.builder() + .addTextContent("Calculation result") + .structuredContent(structuredContent) + .build(); + + return new McpClientTransport() { + Function, Mono> handler; + + @Override + public Mono connect( + Function, Mono> handler) { + this.handler = handler; + return Mono.empty(); + } + + @Override + public Mono closeGracefully() { + return Mono.empty(); + } + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message) { + if (!(message instanceof McpSchema.JSONRPCRequest request)) { + return Mono.empty(); + } + + McpSchema.JSONRPCResponse response; + if (McpSchema.METHOD_INITIALIZE.equals(request.method())) { + response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), MOCK_INIT_RESULT, + null); + } + else if (McpSchema.METHOD_TOOLS_LIST.equals(request.method())) { + response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), mockToolsResult, + null); + } + else if (McpSchema.METHOD_TOOLS_CALL.equals(request.method())) { + response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), + mockCallToolResult, null); + } + else { + return Mono.empty(); + } + + return handler.apply(Mono.just(response)).then(); + } + + @Override + public T unmarshalFrom(Object data, TypeRef typeRef) { + return JSON_MAPPER.convertValue(data, new TypeRef<>() { + @Override + public java.lang.reflect.Type getType() { + return typeRef.getType(); + } + }); + } + }; + } + @Test void validateContextPassedToTransportConnect() { McpClientTransport transport = new McpClientTransport() { @@ -87,4 +181,64 @@ class McpAsyncClientTests { }).doesNotThrowAnyException(); } + @Test + void testCallToolWithOutputSchemaValidationSuccess() throws JsonProcessingException { + McpClientTransport transport = createMockTransportForToolValidation(true, false); + + McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build(); + + StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete(); + + StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3")))) + .expectNextMatches(response -> { + assertThat(response).isNotNull(); + assertThat(response.isError()).isFalse(); + assertThat(response.structuredContent()).isInstanceOf(Map.class); + assertThat((Map) response.structuredContent()).hasSize(2); + assertThat(response.content()).hasSize(1); + return true; + }) + .verifyComplete(); + + StepVerifier.create(client.closeGracefully()).verifyComplete(); + } + + @Test + void testCallToolWithNoOutputSchemaSuccess() throws JsonProcessingException { + McpClientTransport transport = createMockTransportForToolValidation(false, false); + + McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build(); + + StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete(); + + StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3")))) + .expectNextMatches(response -> { + assertThat(response).isNotNull(); + assertThat(response.isError()).isFalse(); + assertThat(response.structuredContent()).isInstanceOf(Map.class); + assertThat((Map) response.structuredContent()).hasSize(2); + assertThat(response.content()).hasSize(1); + return true; + }) + .verifyComplete(); + + StepVerifier.create(client.closeGracefully()).verifyComplete(); + } + + @Test + void testCallToolWithOutputSchemaValidationFailure() throws JsonProcessingException { + McpClientTransport transport = createMockTransportForToolValidation(true, true); + + McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build(); + + StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete(); + + StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3")))) + .expectErrorMatches(ex -> ex instanceof IllegalArgumentException + && ex.getMessage().contains("Tool call result validation failed")) + .verify(); + + StepVerifier.create(client.closeGracefully()).verifyComplete(); + } + } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java index 3feb1d0..a94b9b6 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java @@ -8,9 +8,9 @@ import java.time.Duration; import java.util.List; import io.modelcontextprotocol.MockMcpClientTransport; -import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.InitializeResult; +import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -22,7 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ class McpClientProtocolVersionTests { - private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(30); + private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(300); private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0"); @@ -46,13 +46,12 @@ class McpClientProtocolVersionTests { assertThat(initRequest.protocolVersion()).isEqualTo(transport.protocolVersions().get(0)); transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), - new McpSchema.InitializeResult(protocolVersion, null, + new McpSchema.InitializeResult(protocolVersion, ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), null), null)); }).assertNext(result -> { assertThat(result.protocolVersion()).isEqualTo(protocolVersion); }).verifyComplete(); - } finally { // Ensure cleanup happens even if test fails @@ -81,7 +80,7 @@ class McpClientProtocolVersionTests { assertThat(initRequest.protocolVersion()).isIn(List.of(oldVersion, McpSchema.LATEST_PROTOCOL_VERSION)); transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), - new McpSchema.InitializeResult(oldVersion, null, + new McpSchema.InitializeResult(oldVersion, ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), null), null)); }).assertNext(result -> { @@ -110,7 +109,7 @@ class McpClientProtocolVersionTests { assertThat(request.params()).isInstanceOf(McpSchema.InitializeRequest.class); transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), - new McpSchema.InitializeResult(unsupportedVersion, null, + new McpSchema.InitializeResult(unsupportedVersion, ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), null), null)); }).expectError(RuntimeException.class).verify(); @@ -143,7 +142,7 @@ class McpClientProtocolVersionTests { assertThat(initRequest.protocolVersion()).isEqualTo(latestVersion); transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), - new McpSchema.InitializeResult(latestVersion, null, + new McpSchema.InitializeResult(latestVersion, ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), null), null)); }).assertNext(result -> {