From 36b951558d7044c5b4d5366cabd4c165b87982de Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sun, 29 Dec 2024 17:53:40 +0100 Subject: [PATCH] refactor: improve MCP client code organization and robustness - Add null parameter validation in client builders - Enhance JavaDoc documentation for tools and resources - Reorganize McpAsyncClient methods by functionality - Fix sampling handler test setup --- .../ai/mcp/client/McpAsyncClient.java | 334 ++++++++++-------- .../ai/mcp/client/McpClient.java | 8 + .../ai/mcp/client/McpSyncClient.java | 26 +- .../ai/mcp/spec/McpSchema.java | 48 ++- .../ai/mcp/spec/McpTransport.java | 7 + .../client/AbstractMcpAsyncClientTests.java | 4 +- .../McpAsyncClientResponseHandlerTests.java | 1 - 7 files changed, 268 insertions(+), 160 deletions(-) diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java b/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java index fd39051..5b50266 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java @@ -47,6 +47,7 @@ import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult; import org.springframework.ai.mcp.spec.McpSchema.PaginatedRequest; import org.springframework.ai.mcp.spec.McpSchema.Root; import org.springframework.ai.mcp.spec.McpTransport; +import org.springframework.ai.mcp.util.Assert; import org.springframework.ai.mcp.util.Utils; /** @@ -119,6 +120,10 @@ public class McpAsyncClient { List>> promptsChangeConsumers, Function samplingHandler) { + Assert.notNull(transport, "Transport must not be null"); + Assert.notNull(requestTimeout, "Request timeout must not be null"); + Assert.notNull(clientInfo, "Client info must not be null"); + this.clientInfo = clientInfo; this.clientCapabilities = (clientCapabilities != null) ? clientCapabilities @@ -177,133 +182,9 @@ public class McpAsyncClient { } - public Mono addRoot(Root root) { - - if (root == null) { - return Mono.error(new McpError("Root must not be null")); - } - - if (this.clientCapabilities.roots() == null) { - return Mono.error(new McpError("Client must be configured with roots capabilities")); - } - - if (this.roots.containsKey(root.uri())) { - return Mono.error(new McpError("Root with uri '" + root.uri() + "' already exists")); - } - - this.roots.put(root.uri(), root); - - logger.info("Added root: {}", root); - - if (this.clientCapabilities.roots().listChanged()) { - return this.rootsListChangedNotification(); - } - return Mono.empty(); - } - - public Mono removeRoot(String rootUri) { - - if (rootUri == null) { - return Mono.error(new McpError("Root uri must not be null")); - } - - if (this.clientCapabilities.roots() == null) { - return Mono.error(new McpError("Client must be configured with roots capabilities")); - } - - Root removed = this.roots.remove(rootUri); - - if (removed != null) { - logger.info("Removed Root: {}", rootUri); - if (this.clientCapabilities.roots().listChanged()) { - return this.rootsListChangedNotification(); - } - return Mono.empty(); - } - return Mono.error(new McpError("Root with uri '" + rootUri + "' not found")); - } - - private RequestHandler samplingCreateMessageHandler() { - return params -> { - McpSchema.CreateMessageRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); - - CreateMessageResult response = this.samplingHandler.apply(request); - - return Mono.just(response); - }; - } - - private RequestHandler rootsListRequestHandler() { - return params -> { - McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); - - List roots = this.roots.values().stream().toList(); - - return Mono.just(roots); - }; - } - - private NotificationHandler toolsChangeNotificationHandler( - List>> toolsChangeConsumers) { - - return new NotificationHandler() { - @Override - public Mono handle(Object params) { - // TODO: add support for cursor/pagination - return listTools().flatMap(listToolsResult -> Mono.fromRunnable(() -> { - for (Consumer> toolsChangeConsumer : toolsChangeConsumers) { - toolsChangeConsumer.accept(listToolsResult.tools()); - } - }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { - logger.error("Error handling tools list change notification", error); - return Mono.empty(); - }).then(); // Convert to Mono - } - }; - }; - - private NotificationHandler resourcesChangeNotificationHandler( - List>> resourcesChangeConsumers) { - - return new NotificationHandler() { - @Override - public Mono handle(Object params) { - // TODO: add support for cursor/pagination - return listResources().flatMap(listResourcesResult -> Mono.fromRunnable(() -> { - for (Consumer> resourceChangeConsumer : resourcesChangeConsumers) { - resourceChangeConsumer.accept(listResourcesResult.resources()); - } - }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { - logger.error("Error handling resources list change notification", error); - return Mono.empty(); - }).then(); // Convert to Mono - } - }; - }; - - private NotificationHandler promptsChangeNotificationHandler( - List>> promptsChangeConsumers) { - - return new NotificationHandler() { - @Override - public Mono handle(Object params) { - // TODO: add support for cursor/pagination - return listPrompts().flatMap(listPromptsResult -> Mono.fromRunnable(() -> { - for (Consumer> promptChangeConsumer : promptsChangeConsumers) { - promptChangeConsumer.accept(listPromptsResult.prompts()); - } - }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { - logger.error("Error handling prompts list change notification", error); - return Mono.empty(); - }).then(); // Convert to Mono - } - }; - }; - + // -------------------------- + // Lifecycle + // -------------------------- /** * The initialization phase MUST be the first interaction between client and server. * During this phase, the client and server: @@ -358,13 +239,18 @@ public class McpAsyncClient { }); } - /** - * Send a roots/list_changed notification. - */ - public Mono rootsListChangedNotification() { - return this.mcpSession.sendNotification("notifications/roots/list_changed"); + public void close() { + this.mcpSession.close(); } + public Mono closeGracefully() { + return this.mcpSession.closeGracefully(); + } + + // -------------------------- + // Basic Utilites + // -------------------------- + /** * Send a synchronous ping request. */ @@ -373,6 +259,90 @@ public class McpAsyncClient { }); } + // -------------------------- + // Roots + // -------------------------- + public Mono addRoot(Root root) { + + if (root == null) { + return Mono.error(new McpError("Root must not be null")); + } + + if (this.clientCapabilities.roots() == null) { + return Mono.error(new McpError("Client must be configured with roots capabilities")); + } + + if (this.roots.containsKey(root.uri())) { + return Mono.error(new McpError("Root with uri '" + root.uri() + "' already exists")); + } + + this.roots.put(root.uri(), root); + + logger.info("Added root: {}", root); + + if (this.clientCapabilities.roots().listChanged()) { + return this.rootsListChangedNotification(); + } + return Mono.empty(); + } + + public Mono removeRoot(String rootUri) { + + if (rootUri == null) { + return Mono.error(new McpError("Root uri must not be null")); + } + + if (this.clientCapabilities.roots() == null) { + return Mono.error(new McpError("Client must be configured with roots capabilities")); + } + + Root removed = this.roots.remove(rootUri); + + if (removed != null) { + logger.info("Removed Root: {}", rootUri); + if (this.clientCapabilities.roots().listChanged()) { + return this.rootsListChangedNotification(); + } + return Mono.empty(); + } + return Mono.error(new McpError("Root with uri '" + rootUri + "' not found")); + } + + /** + * Manually, send a roots/list_changed notification. The addRoot and removeRoot + * methods automatically send the roots/list_changed notification. + */ + public Mono rootsListChangedNotification() { + return this.mcpSession.sendNotification("notifications/roots/list_changed"); + } + + private RequestHandler rootsListRequestHandler() { + return params -> { + McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, + new TypeReference() { + }); + + List roots = this.roots.values().stream().toList(); + + return Mono.just(roots); + }; + } + + // -------------------------- + // Sampling + // -------------------------- + private RequestHandler samplingCreateMessageHandler() { + return params -> { + McpSchema.CreateMessageRequest request = transport.unmarshalFrom(params, + new TypeReference() { + }); + + CreateMessageResult response = this.samplingHandler.apply(request); + + return Mono.just(response); + }; + } + // -------------------------- // Tools // -------------------------- @@ -383,32 +353,74 @@ public class McpAsyncClient { }; /** - * Send a tools/call request. - * @param callToolRequest the call tool request. - * @return the call tool result. + * Calls a tool provided by the server. Tools enable servers to expose executable + * functionality that can interact with external systems, perform computations, and + * take actions in the real world. + * @param callToolRequest The request containing: - name: The name of the tool to call + * (must match a tool name from tools/list) - arguments: Arguments that conform to the + * tool's input schema + * @return A Mono that emits the tool execution result containing: - content: List of + * content items (text, images, or embedded resources) representing the tool's output + * - isError: Boolean indicating if the execution failed (true) or succeeded + * (false/absent) */ public Mono callTool(McpSchema.CallToolRequest callToolRequest) { return this.mcpSession.sendRequest("tools/call", callToolRequest, CALL_TOOL_RESULT_TYPE_REF); } /** - * Send a tools/list request. - * @return the list of tools result. + * Retrieves the list of all tools provided by the server. + * @return A Mono that emits the list of tools result containing: - tools: List of + * available tools, each with a name, description, and input schema - nextCursor: + * Optional cursor for pagination if more tools are available */ public Mono listTools() { return this.listTools(null); } /** - * Send a tools/list request. - * @param cursor the cursor - * @return the list of tools result. + * Retrieves a paginated list of tools provided by the server. + * @param cursor Optional pagination cursor from a previous list request + * @return A Mono that emits the list of tools result containing: - tools: List of + * available tools, each with a name, description, and input schema - nextCursor: + * Optional cursor for pagination if more tools are available */ public Mono listTools(String cursor) { return this.mcpSession.sendRequest("tools/list", new McpSchema.PaginatedRequest(cursor), LIST_TOOLS_RESULT_TYPE_REF); } + /** + * Creates a notification handler for tools/list_changed notifications from the + * server. When the server's available tools change, it sends a notification to inform + * connected clients. This handler automatically fetches the updated tool list and + * distributes it to all registered consumers. + * @param toolsChangeConsumers List of consumers that will be notified when the tools + * list changes. Each consumer receives the complete updated list of tools. + * @return A NotificationHandler that processes tools/list_changed notifications by: + * 1. Fetching the current list of tools from the server 2. Distributing the updated + * list to all registered consumers 3. Handling any errors that occur during this + * process + */ + private NotificationHandler toolsChangeNotificationHandler( + List>> toolsChangeConsumers) { + + return new NotificationHandler() { + @Override + public Mono handle(Object params) { + // TODO: add support for cursor/pagination + return listTools().flatMap(listToolsResult -> Mono.fromRunnable(() -> { + for (Consumer> toolsChangeConsumer : toolsChangeConsumers) { + toolsChangeConsumer.accept(listToolsResult.tools()); + } + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling tools list change notification", error); + return Mono.empty(); + }).then(); // Convert to Mono + } + }; + }; + // -------------------------- // Resources // -------------------------- @@ -512,6 +524,25 @@ public class McpAsyncClient { return this.mcpSession.sendRequest("resources/unsubscribe", unsubscribeRequest, VOID_TYPE_REFERENCE); } + private NotificationHandler resourcesChangeNotificationHandler( + List>> resourcesChangeConsumers) { + + return new NotificationHandler() { + @Override + public Mono handle(Object params) { + // TODO: add support for cursor/pagination + return listResources().flatMap(listResourcesResult -> Mono.fromRunnable(() -> { + for (Consumer> resourceChangeConsumer : resourcesChangeConsumers) { + resourceChangeConsumer.accept(listResourcesResult.resources()); + } + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling resources list change notification", error); + return Mono.empty(); + }).then(); // Convert to Mono + } + }; + }; + // -------------------------- // Prompts // -------------------------- @@ -556,12 +587,23 @@ public class McpAsyncClient { return this.mcpSession.sendNotification("notifications/prompts/list_changed"); } - public void close() { - this.mcpSession.close(); - } + private NotificationHandler promptsChangeNotificationHandler( + List>> promptsChangeConsumers) { - public Mono closeGracefully() { - return this.mcpSession.closeGracefully(); - } + return new NotificationHandler() { + @Override + public Mono handle(Object params) { + // TODO: add support for cursor/pagination + return listPrompts().flatMap(listPromptsResult -> Mono.fromRunnable(() -> { + for (Consumer> promptChangeConsumer : promptsChangeConsumers) { + promptChangeConsumer.accept(listPromptsResult.prompts()); + } + }).subscribeOn(Schedulers.boundedElastic())).onErrorResume(error -> { + logger.error("Error handling prompts list change notification", error); + return Mono.empty(); + }).then(); // Convert to Mono + } + }; + }; } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/McpClient.java b/mcp/src/main/java/org/springframework/ai/mcp/client/McpClient.java index f0ef753..3c4663a 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/McpClient.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/McpClient.java @@ -108,16 +108,19 @@ public interface McpClient { } public Builder capabilities(ClientCapabilities capabilities) { + Assert.notNull(capabilities, "Capabilities must not be null"); this.capabilities = capabilities; return this; } public Builder clientInfo(Implementation clientInfo) { + Assert.notNull(clientInfo, "Client info must not be null"); this.clientInfo = clientInfo; return this; } public Builder roots(List roots) { + Assert.notNull(roots, "Roots must not be null"); for (Root root : roots) { this.roots.put(root.uri(), root); } @@ -125,6 +128,7 @@ public interface McpClient { } public Builder roots(Root... roots) { + Assert.notNull(roots, "Roots must not be null"); for (Root root : roots) { this.roots.put(root.uri(), root); } @@ -132,21 +136,25 @@ public interface McpClient { } public Builder sampling(Function samplingHandler) { + Assert.notNull(samplingHandler, "Sampling handler must not be null"); this.samplingHandler = samplingHandler; return this; } public Builder toolsChangeConsumer(Consumer> toolsChangeConsumer) { + Assert.notNull(toolsChangeConsumer, "Tools change consumer must not be null"); this.toolsChangeConsumers.add(toolsChangeConsumer); return this; } public Builder resourcesChangeConsumer(Consumer> resourcesChangeConsumer) { + Assert.notNull(resourcesChangeConsumer, "Resources change consumer must not be null"); this.resourcesChangeConsumers.add(resourcesChangeConsumer); return this; } public Builder promptsChangeConsumer(Consumer> promptsChangeConsumer) { + Assert.notNull(promptsChangeConsumer, "Prompts change consumer must not be null"); this.promptsChangeConsumers.add(promptsChangeConsumer); return this; } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java b/mcp/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java index 18daf13..54fed7b 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java @@ -130,26 +130,36 @@ public class McpSyncClient implements AutoCloseable { // Tools // -------------------------- /** - * Send a tools/call request. - * @param callToolRequest the call tool request. - * @return the call tool result. + * Calls a tool provided by the server. Tools enable servers to expose executable + * functionality that can interact with external systems, perform computations, and + * take actions in the real world. + * @param callToolRequest The request containing: - name: The name of the tool to call + * (must match a tool name from tools/list) - arguments: Arguments that conform to the + * tool's input schema + * @return The tool execution result containing: - content: List of content items + * (text, images, or embedded resources) representing the tool's output - isError: + * Boolean indicating if the execution failed (true) or succeeded (false/absent) */ public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) { return this.delegate.callTool(callToolRequest).block(); } /** - * Send a tools/list request. - * @return the list of tools result. + * Retrieves the list of all tools provided by the server. + * @return The list of tools result containing: - tools: List of available tools, each + * with a name, description, and input schema - nextCursor: Optional cursor for + * pagination if more tools are available */ public McpSchema.ListToolsResult listTools() { return this.delegate.listTools().block(); } /** - * Send a tools/list request. - * @param cursor the cursor - * @return the list of tools result. + * Retrieves a paginated list of tools provided by the server. + * @param cursor Optional pagination cursor from a previous list request + * @return The list of tools result containing: - tools: List of available tools, each + * with a name, description, and input schema - nextCursor: Optional cursor for + * pagination if more tools are available */ public McpSchema.ListToolsResult listTools(String cursor) { return this.delegate.listTools(cursor).block(); diff --git a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java index 7d261c3..ece13f6 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java @@ -364,6 +364,8 @@ public class McpSchema { * by clients to improve the LLM's understanding of available resources. It can be * thought of like a "hint" to the model. * @param mimeType The MIME type of this resource, if known. + * @param annotations Optional annotations for the client. The client can use + * annotations to inform how objects are used or displayed. */ @JsonInclude(JsonInclude.Include.NON_ABSENT) @JsonIgnoreProperties(ignoreUnknown = true) @@ -387,6 +389,8 @@ public class McpSchema { * by clients to improve the LLM's understanding of available resources. It can be * thought of like a "hint" to the model. * @param mimeType The MIME type of this resource, if known. + * @param annotations Optional annotations for the client. The client can use + * annotations to inform how objects are used or displayed. * @see RFC 6570 */ @JsonInclude(JsonInclude.Include.NON_ABSENT) @@ -574,12 +578,32 @@ public class McpSchema { // --------------------------- // Tool Interfaces // --------------------------- + /** + * The server's response to a tools/list request from the client. + * + * @param tools A list of tools that the server provides. + * @param nextCursor An optional cursor for pagination. If present, indicates there + * are more tools available. + */ @JsonInclude(JsonInclude.Include.NON_ABSENT) public record ListToolsResult( // @formatter:off @JsonProperty("tools") List tools, @JsonProperty("nextCursor") String nextCursor) { }// @formatter:on + /** + * Represents a tool that the server provides. Tools enable servers to expose + * executable functionality to the system. Through these tools, you can interact with + * external systems, perform computations, and take actions in the real world. + * + * @param name A unique identifier for the tool. This name is used when calling the + * tool. + * @param description A human-readable description of what the tool does. This can be + * used by clients to improve the LLM's understanding of available tools. + * @param inputSchema A JSON Schema object that describes the expected structure of + * the arguments when calling this tool. This allows clients to validate tool + * arguments before sending them to the server. + */ @JsonInclude(JsonInclude.Include.NON_ABSENT) public record Tool( // @formatter:off @JsonProperty("name") String name, @@ -601,12 +625,28 @@ public class McpSchema { } } + /** + * Used by the client to call a tool provided by the server. + * + * @param name The name of the tool to call. This must match a tool name from + * tools/list. + * @param arguments Arguments to pass to the tool. These must conform to the tool's + * input schema. + */ @JsonInclude(JsonInclude.Include.NON_ABSENT) public record CallToolRequest(// @formatter:off @JsonProperty("name") String name, @JsonProperty("arguments") Map arguments) implements Request { }// @formatter:off + /** + * The server's response to a tools/call request from the client. + * + * @param content A list of content items representing the tool's output. Each item can be text, an image, + * or an embedded resource. + * @param isError If true, indicates that the tool execution failed and the content contains error information. + * If false or absent, indicates successful execution. + */ @JsonInclude(JsonInclude.Include.NON_ABSENT) public record CallToolResult( // @formatter:off @JsonProperty("content") List content, @@ -643,16 +683,16 @@ public class McpSchema { @JsonInclude(JsonInclude.Include.NON_ABSENT) public record CreateMessageResult(// @formatter:off - @JsonProperty("role") Role role, - @JsonProperty("content") Content content, - @JsonProperty("model") String model, + @JsonProperty("role") Role role, + @JsonProperty("content") Content content, + @JsonProperty("model") String model, @JsonProperty("stopReason") StopReason stopReason) { public enum StopReason { @JsonProperty("end_turn") END_TURN, @JsonProperty("stop_sequence") STOP_SEQUENCE, @JsonProperty("max_tokens") MAX_TOKENS - } + } }// @formatter:on // --------------------------- diff --git a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java index f81a6c6..48c6965 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java @@ -93,6 +93,13 @@ public interface McpTransport { */ Mono sendMessage(JSONRPCMessage message); + /** + * Unmarshals the given data into an object of the specified type. + * @param the type of the object to unmarshal + * @param data the data to unmarshal + * @param typeRef the type reference for the object to unmarshal + * @return the unmarshalled object + */ T unmarshalFrom(Object data, TypeReference typeRef); } diff --git a/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java b/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java index 72c683a..b7b26a2 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java @@ -22,6 +22,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import reactor.test.StepVerifier; @@ -225,7 +226,8 @@ public abstract class AbstractMcpAsyncClientTests { .hasMessageContaining("Root with uri 'nonexistent-uri' not found"); } - // @Test + @Test + @Disabled void testReadResource() { StepVerifier.create(mcpAsyncClient.listResources()).consumeNextWith(resources -> { if (!resources.resources().isEmpty()) { diff --git a/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java index d13a8bd..282744c 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java @@ -365,7 +365,6 @@ class McpAsyncClientResponseHandlerTests { // Create client with sampling capability but null handler McpAsyncClient asyncMcpClient = McpClient.using(transport) .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(null) // Null handler .async(); // Create a mock create message request