From c39f5158dd23c78fa4c8af2c051dc3012a962d6b Mon Sep 17 00:00:00 2001 From: lance Date: Thu, 23 Oct 2025 11:22:26 +0800 Subject: [PATCH 1/2] Fix: prevent infinite recursion in listTools() when nextCursor is empty string Signed-off-by: lance --- .../client/McpAsyncClient.java | 5 +- .../client/McpAsyncClientTests.java | 488 ++++++++++-------- 2 files changed, 275 insertions(+), 218 deletions(-) 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 53a05aec3..7c28b4286 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -626,7 +626,10 @@ private McpSchema.CallToolResult validateToolResult(String toolName, McpSchema.C */ public Mono listTools() { return this.listTools(McpSchema.FIRST_PAGE) - .expand(result -> (result.nextCursor() != null) ? this.listTools(result.nextCursor()) : Mono.empty()) + .expand(result -> { + String next = result.nextCursor(); + return (next != null && !next.isEmpty()) ? this.listTools(next) : Mono.empty(); + }) .reduce(new McpSchema.ListToolsResult(new ArrayList<>(), null), (allToolsResult, result) -> { allToolsResult.tools().addAll(result.tools()); return allToolsResult; 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 970d8f257..8eb69cadb 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java @@ -8,19 +8,17 @@ import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.ProtocolVersions; - import org.junit.jupiter.api.Test; - -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.Set; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; +import java.util.stream.Collectors; import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER; import static org.assertj.core.api.Assertions.assertThat; @@ -28,217 +26,273 @@ class McpAsyncClientTests { - public static final McpSchema.Implementation MOCK_SERVER_INFO = new McpSchema.Implementation("test-server", - "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( - ProtocolVersions.MCP_2024_11_05, MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions"); - - 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() { - Function, Mono> handler; - - final AtomicReference contextValue = new AtomicReference<>(); - - @Override - public Mono connect( - Function, Mono> handler) { - return Mono.deferContextual(ctx -> { - this.handler = handler; - if (ctx.hasKey(CONTEXT_KEY)) { - this.contextValue.set(ctx.get(CONTEXT_KEY)); - } - return Mono.empty(); - }); - } - - @Override - public Mono closeGracefully() { - return Mono.empty(); - } - - @Override - public Mono sendMessage(McpSchema.JSONRPCMessage message) { - if (!"hello".equals(this.contextValue.get())) { - return Mono.error(new RuntimeException("Context value not propagated via #connect method")); - } - // We're only interested in handling the init request to provide an init - // response - if (!(message instanceof McpSchema.JSONRPCRequest)) { - return Mono.empty(); - } - McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, - ((McpSchema.JSONRPCRequest) message).id(), MOCK_INIT_RESULT, null); - return handler.apply(Mono.just(initResponse)).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(); - } - }); - } - }; - - assertThatCode(() -> { - McpAsyncClient client = McpClient.async(transport).build(); - client.initialize().contextWrite(ctx -> ctx.put(CONTEXT_KEY, "hello")).block(); - }).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(); - } - + public static final McpSchema.Implementation MOCK_SERVER_INFO = new McpSchema.Implementation("test-server", + "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( + ProtocolVersions.MCP_2024_11_05, MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions"); + + private static final String CONTEXT_KEY = "context.key"; + + private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput) { + + // 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() { + Function, Mono> handler; + + final AtomicReference contextValue = new AtomicReference<>(); + + @Override + public Mono connect( + Function, Mono> handler) { + return Mono.deferContextual(ctx -> { + this.handler = handler; + if (ctx.hasKey(CONTEXT_KEY)) { + this.contextValue.set(ctx.get(CONTEXT_KEY)); + } + return Mono.empty(); + }); + } + + @Override + public Mono closeGracefully() { + return Mono.empty(); + } + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message) { + if (!"hello".equals(this.contextValue.get())) { + return Mono.error(new RuntimeException("Context value not propagated via #connect method")); + } + // We're only interested in handling the init request to provide an init + // response + if (!(message instanceof McpSchema.JSONRPCRequest)) { + return Mono.empty(); + } + McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + ((McpSchema.JSONRPCRequest) message).id(), MOCK_INIT_RESULT, null); + return handler.apply(Mono.just(initResponse)).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(); + } + }); + } + }; + + assertThatCode(() -> { + McpAsyncClient client = McpClient.async(transport).build(); + client.initialize().contextWrite(ctx -> ctx.put(CONTEXT_KEY, "hello")).block(); + }).doesNotThrowAnyException(); + } + + @Test + void testCallToolWithOutputSchemaValidationSuccess() { + 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() { + 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() { + 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(); + } + + @Test + void testListToolsWithEmptyCursor() { + McpSchema.Tool addTool = McpSchema.Tool.builder().name("add").description("calculate add").build(); + McpSchema.Tool subtractTool = McpSchema.Tool.builder().name("subtract").description("calculate subtract").build(); + McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(addTool, subtractTool), ""); + + McpClientTransport transport = new McpClientTransport() { + Function, Mono> handler; + + @Override + public Mono connect(Function, Mono> handler) { + return Mono.deferContextual(ctx -> { + 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 { + 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(); + } + }); + } + }; + + McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build(); + + Mono mono = client.listTools(); + McpSchema.ListToolsResult toolsResult = mono.block(); + assertThat(toolsResult).isNotNull(); + + Set names = toolsResult.tools().stream().map(McpSchema.Tool::name).collect(Collectors.toSet()); + assertThat(names).containsExactlyInAnyOrder("subtract", "add"); + } } From e6559ca757a5b222a8925c280f027d310aa0a0d0 Mon Sep 17 00:00:00 2001 From: lance Date: Tue, 28 Oct 2025 21:43:08 +0800 Subject: [PATCH 2/2] fix McpAsyncClient#listTools prevent infinite recursion Signed-off-by: lance --- .../client/McpAsyncClient.java | 19 +- .../client/McpAsyncClientTests.java | 566 +++++++++--------- 2 files changed, 298 insertions(+), 287 deletions(-) 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 7c28b4286..2d1f4b43c 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -402,6 +402,7 @@ public Mono closeGracefully() { // -------------------------- // Initialization // -------------------------- + /** * The initialization phase should be the first interaction between client and server. * The client will ensure it happens in case it has not been explicitly called and in @@ -448,6 +449,7 @@ public Mono ping() { // -------------------------- // Roots // -------------------------- + /** * Adds a new root to the client's root list. * @param root The root to add. @@ -625,16 +627,13 @@ private McpSchema.CallToolResult validateToolResult(String toolName, McpSchema.C * @return A Mono that emits the list of all tools result */ public Mono listTools() { - return this.listTools(McpSchema.FIRST_PAGE) - .expand(result -> { - String next = result.nextCursor(); - return (next != null && !next.isEmpty()) ? this.listTools(next) : Mono.empty(); - }) - .reduce(new McpSchema.ListToolsResult(new ArrayList<>(), null), (allToolsResult, result) -> { - allToolsResult.tools().addAll(result.tools()); - return allToolsResult; - }) - .map(result -> new McpSchema.ListToolsResult(Collections.unmodifiableList(result.tools()), null)); + return this.listTools(McpSchema.FIRST_PAGE).expand(result -> { + String next = result.nextCursor(); + return (next != null && !next.isEmpty()) ? this.listTools(next) : Mono.empty(); + }).reduce(new McpSchema.ListToolsResult(new ArrayList<>(), null), (allToolsResult, result) -> { + allToolsResult.tools().addAll(result.tools()); + return allToolsResult; + }).map(result -> new McpSchema.ListToolsResult(Collections.unmodifiableList(result.tools()), null)); } /** 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 8eb69cadb..48bf1da5b 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java @@ -4,14 +4,6 @@ package io.modelcontextprotocol.client; -import io.modelcontextprotocol.json.TypeRef; -import io.modelcontextprotocol.spec.McpClientTransport; -import io.modelcontextprotocol.spec.McpSchema; -import io.modelcontextprotocol.spec.ProtocolVersions; -import org.junit.jupiter.api.Test; -import reactor.core.publisher.Mono; -import reactor.test.StepVerifier; - import java.util.List; import java.util.Map; import java.util.Objects; @@ -20,279 +12,299 @@ import java.util.function.Function; import java.util.stream.Collectors; +import io.modelcontextprotocol.json.TypeRef; +import io.modelcontextprotocol.spec.McpClientTransport; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.ProtocolVersions; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + 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 { - public static final McpSchema.Implementation MOCK_SERVER_INFO = new McpSchema.Implementation("test-server", - "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( - ProtocolVersions.MCP_2024_11_05, MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions"); - - private static final String CONTEXT_KEY = "context.key"; - - private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput) { - - // 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() { - Function, Mono> handler; - - final AtomicReference contextValue = new AtomicReference<>(); - - @Override - public Mono connect( - Function, Mono> handler) { - return Mono.deferContextual(ctx -> { - this.handler = handler; - if (ctx.hasKey(CONTEXT_KEY)) { - this.contextValue.set(ctx.get(CONTEXT_KEY)); - } - return Mono.empty(); - }); - } - - @Override - public Mono closeGracefully() { - return Mono.empty(); - } - - @Override - public Mono sendMessage(McpSchema.JSONRPCMessage message) { - if (!"hello".equals(this.contextValue.get())) { - return Mono.error(new RuntimeException("Context value not propagated via #connect method")); - } - // We're only interested in handling the init request to provide an init - // response - if (!(message instanceof McpSchema.JSONRPCRequest)) { - return Mono.empty(); - } - McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, - ((McpSchema.JSONRPCRequest) message).id(), MOCK_INIT_RESULT, null); - return handler.apply(Mono.just(initResponse)).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(); - } - }); - } - }; - - assertThatCode(() -> { - McpAsyncClient client = McpClient.async(transport).build(); - client.initialize().contextWrite(ctx -> ctx.put(CONTEXT_KEY, "hello")).block(); - }).doesNotThrowAnyException(); - } - - @Test - void testCallToolWithOutputSchemaValidationSuccess() { - 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() { - 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() { - 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(); - } - - @Test - void testListToolsWithEmptyCursor() { - McpSchema.Tool addTool = McpSchema.Tool.builder().name("add").description("calculate add").build(); - McpSchema.Tool subtractTool = McpSchema.Tool.builder().name("subtract").description("calculate subtract").build(); - McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(addTool, subtractTool), ""); - - McpClientTransport transport = new McpClientTransport() { - Function, Mono> handler; - - @Override - public Mono connect(Function, Mono> handler) { - return Mono.deferContextual(ctx -> { - 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 { - 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(); - } - }); - } - }; - - McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build(); - - Mono mono = client.listTools(); - McpSchema.ListToolsResult toolsResult = mono.block(); - assertThat(toolsResult).isNotNull(); - - Set names = toolsResult.tools().stream().map(McpSchema.Tool::name).collect(Collectors.toSet()); - assertThat(names).containsExactlyInAnyOrder("subtract", "add"); - } + public static final McpSchema.Implementation MOCK_SERVER_INFO = new McpSchema.Implementation("test-server", + "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( + ProtocolVersions.MCP_2024_11_05, MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions"); + + private static final String CONTEXT_KEY = "context.key"; + + private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput) { + + // 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() { + Function, Mono> handler; + + final AtomicReference contextValue = new AtomicReference<>(); + + @Override + public Mono connect( + Function, Mono> handler) { + return Mono.deferContextual(ctx -> { + this.handler = handler; + if (ctx.hasKey(CONTEXT_KEY)) { + this.contextValue.set(ctx.get(CONTEXT_KEY)); + } + return Mono.empty(); + }); + } + + @Override + public Mono closeGracefully() { + return Mono.empty(); + } + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message) { + if (!"hello".equals(this.contextValue.get())) { + return Mono.error(new RuntimeException("Context value not propagated via #connect method")); + } + // We're only interested in handling the init request to provide an init + // response + if (!(message instanceof McpSchema.JSONRPCRequest)) { + return Mono.empty(); + } + McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + ((McpSchema.JSONRPCRequest) message).id(), MOCK_INIT_RESULT, null); + return handler.apply(Mono.just(initResponse)).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(); + } + }); + } + }; + + assertThatCode(() -> { + McpAsyncClient client = McpClient.async(transport).build(); + client.initialize().contextWrite(ctx -> ctx.put(CONTEXT_KEY, "hello")).block(); + }).doesNotThrowAnyException(); + } + + @Test + void testCallToolWithOutputSchemaValidationSuccess() { + 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() { + 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() { + 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(); + } + + @Test + void testListToolsWithEmptyCursor() { + McpSchema.Tool addTool = McpSchema.Tool.builder().name("add").description("calculate add").build(); + McpSchema.Tool subtractTool = McpSchema.Tool.builder() + .name("subtract") + .description("calculate subtract") + .build(); + McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(addTool, subtractTool), ""); + + McpClientTransport transport = new McpClientTransport() { + Function, Mono> handler; + + @Override + public Mono connect( + Function, Mono> handler) { + return Mono.deferContextual(ctx -> { + 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 { + 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(); + } + }); + } + }; + + McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build(); + + Mono mono = client.listTools(); + McpSchema.ListToolsResult toolsResult = mono.block(); + assertThat(toolsResult).isNotNull(); + + Set names = toolsResult.tools().stream().map(McpSchema.Tool::name).collect(Collectors.toSet()); + assertThat(names).containsExactlyInAnyOrder("subtract", "add"); + } + }