From 7a98a935af83b4975f2b12955781a3946aa5707d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dariusz=20J=C4=99drzejczyk?= Date: Tue, 17 Dec 2024 18:10:46 +0100 Subject: [PATCH] Refactor the mcp-core module - The core module has been refactored to remove unnecessary abstraction between the SSE transport and the STDIO transport. - Tests have been improved and no longer complain about multiple subscriptions to single-subscriber Sink. - SSE transport layer is now assuming non-blocking execution and no longer makes thread hops when not needed. - Decouple ObjectMapper and error handling from AbstractMcpTransport and remove later. Additional improvements: - Rename transport classes to better reflect client-side role (*ServerTransport -> *ClientTransport) - Introduce fluent builder pattern for client configuration - Simplify DefaultMcpSession by removing ObjectMapper dependency - Move unmarshal responsibility to transport layer - Update documentation and diagrams to reflect new architecture - Clean up test implementations --- README.md | 18 +- spring-ai-mcp-core/README.md | 31 +-- .../docs/CLIENT.SSE.TRANSPORT.md | 6 +- .../docs/CLIENT.STDIO.TRANSPORT.md | 6 +- spring-ai-mcp-core/docs/class-diagrams.puml | 30 +-- .../docs/spring-ai-mcp-uml-classdiagram.svg | 2 +- .../ai/mcp/client/McpAsyncClient.java | 67 +++-- .../ai/mcp/client/McpClient.java | 107 +++++--- .../ai/mcp/client/McpSyncClient.java | 21 +- ...Transport.java => SseClientTransport.java} | 240 +++++++----------- .../ai/mcp/client/stdio/ServerParameters.java | 2 +- ...ansport.java => StdioClientTransport.java} | 162 ++++++++---- .../ai/mcp/spec/AbstractMcpTransport.java | 122 --------- .../ai/mcp/spec/DefaultMcpSession.java | 34 ++- .../ai/mcp/spec/McpTransport.java | 31 +-- .../ai/mcp/{client => }/util/Assert.java | 2 +- .../src/main/resources/logback.xml | 2 +- .../ai/mcp/attic/ClientSessionTests2.java | 9 +- .../ai/mcp/attic/ClientSessionTests3.java | 9 +- .../springframework/ai/mcp/attic/MainSSE.java | 11 +- .../ai/mcp/attic}/ServerParametersParser.java | 4 +- .../client/AbstractMcpAsyncClientTests.java | 19 +- .../client/AbstractMcpSyncClientTests.java | 19 +- ...ests.java => SseClientTransportTests.java} | 52 ++-- .../client/sse/SseMcpAsyncClientTests.java | 7 +- .../mcp/client/sse/SseMcpSyncClientTests.java | 9 +- .../mcp/client/stdio/McpAsyncClientTests.java | 7 +- .../mcp/client/stdio/McpSyncClientTests.java | 27 +- .../ai/mcp/spec/DefaultMcpSessionTests.java | 105 ++++---- .../ai/mcp/spec/DefaultMcpTransportTests.java | 154 ----------- 30 files changed, 557 insertions(+), 758 deletions(-) rename spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/{SseServerTransport.java => SseClientTransport.java} (60%) rename spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/{StdioServerTransport.java => StdioClientTransport.java} (64%) delete mode 100644 spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java rename spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/{client => }/util/Assert.java (98%) rename spring-ai-mcp-core/src/{main/java/org/springframework/ai/mcp/client/stdio => test/java/org/springframework/ai/mcp/attic}/ServerParametersParser.java (94%) rename spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/{SseServerTransportTests.java => SseClientTransportTests.java} (85%) delete mode 100644 spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java diff --git a/README.md b/README.md index 597e2c7..5d8c079 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,10 @@ # Spring AI MCP -Java SDK for the Model Context Protocol (MCP), providing seamless integration between Java and Spring applications and MCP-compliant AI models and tools. +Java SDK for the Model Context Protocol (MCP), providing seamless integration between Java and Spring applications and MCP-compliant AI resources and tools. ## Overview -Spring AI MCP is an experimental project that provides Java and Spring Framework integration for the [Model Context Protocol](https://modelcontextprotocol.org/docs/concepts/architecture). It enables Spring applications to interact with AI models and tools through a standardized interface, supporting both synchronous and asynchronous communication patterns. +Spring AI MCP is an experimental project that provides Java and Spring Framework integration for the [Model Context Protocol](https://modelcontextprotocol.org/docs/concepts/architecture). It enables Java applications to interact with AI models and tools through a standardized interface, supporting both synchronous and asynchronous communication patterns. @@ -12,22 +12,20 @@ Spring AI MCP is an experimental project that provides Java and Spring Framework The project consists of two main modules: -### [spring-ai-mcp-core](./spring-ai-mcp-core/README.md) +### [mcp-core](./spring-ai-mcp-core/README.md) The core module provides a Java implementation of the Model Context Protocol specification. It includes: - Synchronous and asynchronous client implementations -- Standard MCP operations support (tool discovery, resource management, prompt handling) -- Stdio-based server transport -- Reactive programming support using Project Reactor +- Standard MCP operations support (tool discovery, resource management, prompt handling). Support for request and notificaiotn handling. +- [Stdio](https://spec.modelcontextprotocol.io/specification/basic/transports/#stdio) and [SSE](https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse) transport implementations. +- [Find more](./spring-ai-mcp-core/README.md). -[find more](./spring-ai-mcp-core/README.md) - -### [spring-ai-mcp-spring](./spring-ai-mcp-spring/README.md) +### [mcp-spring](./spring-ai-mcp-spring/README.md) The Spring integration module provides Spring-specific functionality: - Integration with Spring AI's function calling system - Spring-friendly abstractions for MCP clients -- Automatic conversion between JSON and Java objects for tool arguments +- Auto-configurations (WIP) ## Requirements diff --git a/spring-ai-mcp-core/README.md b/spring-ai-mcp-core/README.md index 57fecec..c4fb117 100644 --- a/spring-ai-mcp-core/README.md +++ b/spring-ai-mcp-core/README.md @@ -20,7 +20,6 @@ This SDK implements the Model Context Protocol, enabling seamless integration wi - Multiple transport implementations: - Stdio-based transport for process-based communication - SSE-based transport for HTTP streaming -- Reactive programming support using Project Reactor - Configurable request timeouts - Customizable JSON serialization/deserialization @@ -32,43 +31,37 @@ Add the following dependency to your Maven project: org.springframework.experimental spring-ai-mcp-core - 0.1.0 + 0.2.0-SNAPSHOT ``` -## Documentation - -Detailed UML class diagrams showing the relationships between components can be found in [docs/class-diagrams.puml](docs/class-diagrams.puml). The diagrams include: -- Core Components: Shows the main interfaces, classes, and their relationships -- Message Flow: Illustrates the message and resource type hierarchies - ## Usage ### Transport Layer Options The SDK provides two transport implementations: -#### StdioServerTransport +#### StdioClientTransport Standard I/O based transport for process-based communication with MCP servers: ```java ServerParameters params = ServerParameters.builder("npx") .args("-y", "@modelcontextprotocol/server-everything", "dir") .build(); -McpTransport transport = new StdioServerTransport(params); +McpTransport transport = new StdioClientTransport(params); ``` -#### SseServerTransport +#### SseClientTransport Server-Sent Events (SSE) based transport following the MCP HTTP with SSE transport specification: ```java WebClient.Builder webClientBuilder = WebClient.builder() .baseUrl("http://your-mcp-server"); -McpTransport transport = new SseServerTransport(webClientBuilder); +McpTransport transport = new SseClientTransport(webClientBuilder); // Or with custom ObjectMapper ObjectMapper mapper = new ObjectMapper(); -McpTransport transport = new SseServerTransport(webClientBuilder, mapper); +McpTransport transport = new SseClientTransport(webClientBuilder, mapper); ``` The SSE transport provides: @@ -87,7 +80,7 @@ ServerParameters params = ServerParameters.builder("npx") .args("-y", "@modelcontextprotocol/server-everything", "dir") .build(); -try (McpSyncClient client = McpClient.sync(new StdioServerTransport(params))) { +try (McpSyncClient client = McpClient.sync(new StdioClientTransport(params))) { // Initialize connection with protocol version and capabilities McpSchema.InitializeResult initResult = client.initialize(); @@ -123,7 +116,7 @@ ServerParameters params = ServerParameters.builder("npx") // Initialize async client with custom timeout and object mapper McpAsyncClient client = McpClient.async( - new StdioServerTransport(params), + new StdioClientTransport(params), Duration.ofSeconds(30), new ObjectMapper() ); @@ -172,11 +165,11 @@ The SDK follows a layered architecture with clear separation of concerns: - **McpAsyncClient**: Primary async implementation using Project Reactor for non-blocking operations - **McpSyncClient**: Synchronous wrapper around the async client for blocking operations - **McpSession**: Core session interface defining communication patterns -- **McpTransport**: Transport layer interface for server communication +- **McpTransport**: Transport layer interface for client/server communication - **McpSchema**: Comprehensive protocol schema definitions -- **DefaultMcpSession**: Base implementation of the session management -- **StdioServerTransport**: Standard I/O based server communication -- **SseServerTransport**: HTTP-based transport using Server-Sent Events for bidirectional communication +- **DefaultMcpSession**: Default implementation of the session management +- **StdioClientTransport**: Standard I/O based client to server communication +- **SseClientTransport**: HTTP-based transport using Server-Sent Events for bidirectional client-sever communication diff --git a/spring-ai-mcp-core/docs/CLIENT.SSE.TRANSPORT.md b/spring-ai-mcp-core/docs/CLIENT.SSE.TRANSPORT.md index 62bce86..e65794f 100644 --- a/spring-ai-mcp-core/docs/CLIENT.SSE.TRANSPORT.md +++ b/spring-ai-mcp-core/docs/CLIENT.SSE.TRANSPORT.md @@ -2,7 +2,7 @@ ## Overview -The `SseServerTransport` class implements the Model Context Protocol (MCP) HTTP with SSE transport specification, providing a bidirectional communication channel between clients and servers. This implementation is part of the Spring AI MCP Core library and follows the [MCP HTTP with SSE Transport Specification](https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse). +The `SseClientTransport` class implements the Model Context Protocol (MCP) HTTP with SSE transport specification, providing a bidirectional communication channel between clients and servers. This implementation is part of the Spring AI MCP Core library and follows the [MCP HTTP with SSE Transport Specification](https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse). ## Key Features @@ -53,7 +53,7 @@ WebClient.Builder webClientBuilder = WebClient.builder() .baseUrl("http://localhost:3001"); // Initialize transport -SseServerTransport transport = new SseServerTransport(webClientBuilder); +SseClientTransport transport = new SseClientTransport(webClientBuilder); // Create a JSON-RPC request JSONRPCRequest request = new JSONRPCRequest( @@ -141,4 +141,4 @@ The implementation includes comprehensive tests covering: - Shutdown procedures - Edge cases and failure scenarios -For detailed test examples, see `SseServerTransportTests.java`. +For detailed test examples, see `SseClientTransportTests.java`. diff --git a/spring-ai-mcp-core/docs/CLIENT.STDIO.TRANSPORT.md b/spring-ai-mcp-core/docs/CLIENT.STDIO.TRANSPORT.md index b78ab55..1c2dc6f 100644 --- a/spring-ai-mcp-core/docs/CLIENT.STDIO.TRANSPORT.md +++ b/spring-ai-mcp-core/docs/CLIENT.STDIO.TRANSPORT.md @@ -1,6 +1,6 @@ -# StdioServerTransport +# StdioClientTransport -The `StdioServerTransport` class implements the [MCP Stdio Transport](https://spec.modelcontextprotocol.io/specification/basic/transports/#stdio) specification, providing communication with MCP servers over standard input/output streams. +The `StdioClientTransport` class implements the [MCP Stdio Transport](https://spec.modelcontextprotocol.io/specification/basic/transports/#stdio) specification, providing communication with MCP servers over standard input/output streams. ## Overview @@ -84,7 +84,7 @@ ServerParameters params = ServerParameters.builder() .build(); // Create transport -StdioServerTransport transport = new StdioServerTransport(params); +StdioClientTransport transport = new StdioClientTransport(params); // Start transport transport.start(); diff --git a/spring-ai-mcp-core/docs/class-diagrams.puml b/spring-ai-mcp-core/docs/class-diagrams.puml index 14a7f3a..4b2c8cd 100644 --- a/spring-ai-mcp-core/docs/class-diagrams.puml +++ b/spring-ai-mcp-core/docs/class-diagrams.puml @@ -1,25 +1,18 @@ @startuml Core Components interface McpTransport { - +void start() + +Mono connect(Function, + Mono> handler) + +Mono sendMessage() +void close() +Mono closeGracefully() - +void setInboudMessageHandler() - +void setInboundErrorHandler() - +Mono sendMessage() -} - -abstract class AbstractMcpTransport { - #ObjectMapper objectMapper - #getInboundSink() - #getOutboundSink() - #getErrorSink() } interface McpSession { - + Mono sendRequest() - +Mono sendNotification() + + Mono sendRequest(String method, Object requestParams, TypeReference typeRef) + +Mono sendNotification(String method, Map params) +Mono closeGracefully() + +close() } class DefaultMcpSession { @@ -52,12 +45,12 @@ class McpSyncClient { +GetPromptResult getPrompt() } -class StdioServerTransport { +class StdioClientTransport { -Process serverProcess -ServerParameters parameters } -class SseServerTransport { +class SseClientTransport { -WebClient webClient } @@ -74,11 +67,10 @@ class McpError { -String message } -McpTransport <|.. AbstractMcpTransport -AbstractMcpTransport <|-- StdioServerTransport -AbstractMcpTransport <|-- SseServerTransport +McpTransport <|-- StdioClientTransport +McpTransport <|-- SseClientTransport McpSession <|-- DefaultMcpSession -DefaultMcpSession <|-- McpAsyncClient +DefaultMcpSession --o McpAsyncClient McpClient ..> McpAsyncClient : creates McpClient ..> McpSyncClient : creates McpSyncClient o-- McpAsyncClient diff --git a/spring-ai-mcp-core/docs/spring-ai-mcp-uml-classdiagram.svg b/spring-ai-mcp-core/docs/spring-ai-mcp-uml-classdiagram.svg index fc6b72b..6b05be7 100644 --- a/spring-ai-mcp-core/docs/spring-ai-mcp-uml-classdiagram.svg +++ b/spring-ai-mcp-core/docs/spring-ai-mcp-uml-classdiagram.svg @@ -1 +1 @@ -McpTransportvoid start()void close()Mono<Void> closeGracefully()void setInboudMessageHandler()void setInboundErrorHandler()Mono<Void> sendMessage()AbstractMcpTransportObjectMapper objectMappergetInboundSink()getOutboundSink()getErrorSink()McpSession<T> Mono<T> sendRequest()Mono<Void> sendNotification()Mono<Void> closeGracefully()DefaultMcpSessionMcpTransport transportObjectMapper objectMapperMcpClientasync()sync()McpAsyncClientMono<InitializeResult> initialize()Mono<CallToolResult> callTool()Mono<ListToolsResult> listTools()Mono<ListResourcesResult> listResources()Mono<ReadResourceResult> readResource()Mono<ListPromptsResult> listPrompts()Mono<GetPromptResult> getPrompt()McpSyncClientInitializeResult initialize()CallToolResult callTool()ListToolsResult listTools()ListResourcesResult listResources()ReadResourceResult readResource()ListPromptsResult listPrompts()GetPromptResult getPrompt()StdioServerTransportProcess serverProcessServerParameters parametersSseServerTransportWebClient webClientMcpSchemaclass ErrorCodesinterface Requestinterface JSONRPCMessageinterface ResourceContentsinterface ContentMcpErrorint codeString messagecreatescreatesusesthrows \ No newline at end of file +McpTransportMono<Void> connect(Function<Mono<JSONRPCMessage>,Mono<JSONRPCMessage>> handler)Mono<Void> sendMessage()void close()Mono<Void> closeGracefully()McpSession<T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef)Mono<Void> sendNotification(String method, Map<String, Object> params)Mono<Void> closeGracefully()close()DefaultMcpSessionMcpTransport transportObjectMapper objectMapperMcpClientasync()sync()McpAsyncClientMono<InitializeResult> initialize()Mono<CallToolResult> callTool()Mono<ListToolsResult> listTools()Mono<ListResourcesResult> listResources()Mono<ReadResourceResult> readResource()Mono<ListPromptsResult> listPrompts()Mono<GetPromptResult> getPrompt()McpSyncClientInitializeResult initialize()CallToolResult callTool()ListToolsResult listTools()ListResourcesResult listResources()ReadResourceResult readResource()ListPromptsResult listPrompts()GetPromptResult getPrompt()StdioClientTransportProcess serverProcessServerParameters parametersSseClientTransportWebClient webClientMcpSchemaclass ErrorCodesinterface Requestinterface JSONRPCMessageinterface ResourceContentsinterface ContentMcpErrorint codeString messagecreatescreatesusesthrows \ No newline at end of file diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java index ac4cbbc..7accb87 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java @@ -18,7 +18,6 @@ package org.springframework.ai.mcp.client; import java.time.Duration; import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -33,22 +32,41 @@ import org.springframework.ai.mcp.spec.McpSchema.PaginatedRequest; import org.springframework.ai.mcp.spec.McpTransport; /** + * The Model Context Protocol (MCP) client implementation that provides asynchronous + * communication with MCP servers. + * * @author Dariusz Jędrzejczyk * @author Christian Tzolov */ -public class McpAsyncClient extends DefaultMcpSession { +public class McpAsyncClient { private final static Logger logger = LoggerFactory.getLogger(McpAsyncClient.class); private static TypeReference VOID_TYPE_REFERENCE = new TypeReference<>() { }; + /** + * The MCP session implementation that manages bidirectional JSON-RPC communication + * between clients and servers. + */ + private final DefaultMcpSession mcpSession; + + /** + * Create a new McpAsyncClient with the given transport. + * @param transport the transport to use. + */ public McpAsyncClient(McpTransport transport) { - this(transport, Duration.ofSeconds(20), new ObjectMapper()); + this(transport, Duration.ofSeconds(20)); } - public McpAsyncClient(McpTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { - super(requestTimeout, objectMapper, transport); + /** + * 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. + */ + public McpAsyncClient(McpTransport transport, Duration requestTimeout) { + this.mcpSession = new DefaultMcpSession(requestTimeout, transport); } /** @@ -85,7 +103,7 @@ public class McpAsyncClient extends DefaultMcpSession { new McpSchema.ClientCapabilities(null, new McpSchema.ClientCapabilities.RootCapabilities(true), null), new McpSchema.Implementation("mcp-java-client", "0.0.1")); // @formatter:on - Mono result = this.sendRequest("initialize", initializeRequest, + Mono result = this.mcpSession.sendRequest("initialize", initializeRequest, new TypeReference() { }); @@ -100,7 +118,7 @@ public class McpAsyncClient extends DefaultMcpSession { "Unsupported protocol version from the server: " + initializeResult.protocolVersion())); } else { - return this.sendNotification("notifications/initialized", null).thenReturn(initializeResult); + return this.mcpSession.sendNotification("notifications/initialized", null).thenReturn(initializeResult); } }); } @@ -109,14 +127,14 @@ public class McpAsyncClient extends DefaultMcpSession { * Send a roots/list_changed notification. */ public Mono sendRootsListChanged() { - return this.sendNotification("notifications/roots/list_changed"); + return this.mcpSession.sendNotification("notifications/roots/list_changed"); } /** * Send a synchronous ping request. */ public Mono ping() { - return this.sendRequest("ping", null, VOID_TYPE_REFERENCE); + return this.mcpSession.sendRequest("ping", null, VOID_TYPE_REFERENCE); } // -------------------------- @@ -134,7 +152,7 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the call tool result. */ public Mono callTool(McpSchema.CallToolRequest callToolRequest) { - return this.sendRequest("tools/call", callToolRequest, CALL_TOOL_RESULT_TYPE_REF); + return this.mcpSession.sendRequest("tools/call", callToolRequest, CALL_TOOL_RESULT_TYPE_REF); } /** @@ -151,7 +169,8 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the list of tools result. */ public Mono listTools(String cursor) { - return this.sendRequest("tools/list", new McpSchema.PaginatedRequest(cursor), LIST_TOOLS_RESULT_TYPE_REF); + return this.mcpSession.sendRequest("tools/list", new McpSchema.PaginatedRequest(cursor), + LIST_TOOLS_RESULT_TYPE_REF); } // -------------------------- @@ -181,7 +200,7 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the list of resources result. */ public Mono listResources(String cursor) { - return this.sendRequest("resources/list", new McpSchema.PaginatedRequest(cursor), + return this.mcpSession.sendRequest("resources/list", new McpSchema.PaginatedRequest(cursor), LIST_RESOURCES_RESULT_TYPE_REF); } @@ -200,7 +219,7 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the resource content. */ public Mono readResource(McpSchema.ReadResourceRequest readResourceRequest) { - return this.sendRequest("resources/read", readResourceRequest, READ_RESOURCE_RESULT_TYPE_REF); + return this.mcpSession.sendRequest("resources/read", readResourceRequest, READ_RESOURCE_RESULT_TYPE_REF); } /** @@ -223,7 +242,7 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the list of resource templates result. */ public Mono listResourceTemplates(String cursor) { - return this.sendRequest("resources/templates/list", new McpSchema.PaginatedRequest(cursor), + return this.mcpSession.sendRequest("resources/templates/list", new McpSchema.PaginatedRequest(cursor), LIST_RESOURCE_TEMPLATES_RESULT_TYPE_REF); } @@ -232,7 +251,7 @@ public class McpAsyncClient extends DefaultMcpSession { * that declared the listChanged capability SHOULD send a notification: */ public Mono sendResourcesListChanged() { - return this.sendNotification("notifications/resources/list_changed"); + return this.mcpSession.sendNotification("notifications/resources/list_changed"); } /** @@ -245,7 +264,7 @@ public class McpAsyncClient extends DefaultMcpSession { * subscribe to. */ public Mono subscribeResource(McpSchema.SubscribeRequest subscribeRequest) { - return this.sendRequest("resources/subscribe", subscribeRequest, VOID_TYPE_REFERENCE); + return this.mcpSession.sendRequest("resources/subscribe", subscribeRequest, VOID_TYPE_REFERENCE); } /** @@ -254,7 +273,7 @@ public class McpAsyncClient extends DefaultMcpSession { * to unsubscribe from. */ public Mono unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) { - return this.sendRequest("resources/unsubscribe", unsubscribeRequest, VOID_TYPE_REFERENCE); + return this.mcpSession.sendRequest("resources/unsubscribe", unsubscribeRequest, VOID_TYPE_REFERENCE); } // -------------------------- @@ -280,7 +299,7 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the list of prompts result. */ public Mono listPrompts(String cursor) { - return this.sendRequest("prompts/list", new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF); + return this.mcpSession.sendRequest("prompts/list", new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF); } /** @@ -289,7 +308,7 @@ public class McpAsyncClient extends DefaultMcpSession { * @return the get prompt result. */ public Mono getPrompt(GetPromptRequest getPromptRequest) { - return this.sendRequest("prompts/get", getPromptRequest, GET_PROMPT_RESULT_TYPE_REF); + return this.mcpSession.sendRequest("prompts/get", getPromptRequest, GET_PROMPT_RESULT_TYPE_REF); } /** @@ -298,7 +317,15 @@ public class McpAsyncClient extends DefaultMcpSession { * any previous subscription from the client. */ public Mono promptListChangedNotification() { - return this.sendNotification("notifications/prompts/list_changed"); + return this.mcpSession.sendNotification("notifications/prompts/list_changed"); + } + + public void close() { + this.mcpSession.close(); + } + + public Mono closeGracefully() { + return this.mcpSession.closeGracefully(); } } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java index 6f12f96..a9b968b 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpClient.java @@ -18,9 +18,8 @@ package org.springframework.ai.mcp.client; import java.time.Duration; -import com.fasterxml.jackson.databind.ObjectMapper; - import org.springframework.ai.mcp.spec.McpTransport; +import org.springframework.ai.mcp.util.Assert; /** * Factory class providing static methods for creating Model Context Protocol (MCP) @@ -39,8 +38,11 @@ import org.springframework.ai.mcp.spec.McpTransport; * including request timeout and JSON object mapping. * *

- * Future implementations will introduce a builder pattern for more flexible client - * configuration. + * Use the builder pattern for flexible client configuration: 

{@code
+ * McpClient.using(transport)
+ *     .withRequestTimeout(Duration.ofSeconds(5))
+ *     .sync(); // or .async()
+ * }
* * @author Christian Tzolov * @author Dariusz Jędrzejczyk @@ -54,55 +56,88 @@ public class McpClient { private McpClient() { } - // TODO: introduce a builder like: - // McpClient.using(transport) - // .withRequestTimeout(Duration.ofSeconds(5)) - // .withObjectMapper(objectMapper); <- or even a more sophisticated - // JSONtoPOJOCodec type - // .sync(); + /** + * Start building an MCP client with the specified transport. + * @param transport The transport layer implementation for MCP communication + * @return A new builder instance + */ + public static Builder using(McpTransport transport) { + return new Builder(transport); + } /** - * Creates an asynchronous MCP client with default configuration. - * @param transport The transport layer implementation for MCP communication - * @return A new instance of {@link McpAsyncClient} + * Builder class for creating MCP clients with custom configuration. */ + public static class Builder { + + private final McpTransport transport; + + private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout + + private Builder(McpTransport transport) { + Assert.notNull(transport, "Transport must not be null"); + this.transport = transport; + } + + /** + * Set the request timeout duration. + * @param requestTimeout The duration to wait before timing out requests + * @return This builder instance + */ + public Builder withRequestTimeout(Duration requestTimeout) { + Assert.notNull(requestTimeout, "Request timeout must not be null"); + this.requestTimeout = requestTimeout; + return this; + } + + /** + * Build a synchronous MCP client. + * @return A new instance of {@link McpSyncClient} + */ + public McpSyncClient sync() { + return new McpSyncClient(async()); + } + + /** + * Build an asynchronous MCP client. + * @return A new instance of {@link McpAsyncClient} + */ + public McpAsyncClient async() { + return new McpAsyncClient(transport, requestTimeout); + } + + } + + /** + * @deprecated Use {@link #using(McpTransport)} instead. + */ + @Deprecated public static McpAsyncClient async(McpTransport transport) { - return new McpAsyncClient(transport); + return using(transport).async(); } /** - * Creates an asynchronous MCP client with custom configuration. - * @param transport The transport layer implementation for MCP communication - * @param requestTimeout The duration to wait before timing out requests - * @param objectMapper The Jackson object mapper for JSON - * serialization/deserialization - * @return A new instance of {@link McpAsyncClient} + * @deprecated Use {@link #using(McpTransport)} instead. */ - public static McpAsyncClient async(McpTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { - return new McpAsyncClient(transport, requestTimeout, objectMapper); + @Deprecated + public static McpAsyncClient async(McpTransport transport, Duration requestTimeout) { + return using(transport).withRequestTimeout(requestTimeout).async(); } /** - * Creates a synchronous MCP client with default configuration. This method wraps an - * asynchronous client to provide synchronous operations. - * @param transport The transport layer implementation for MCP communication - * @return A new instance of {@link McpSyncClient} + * @deprecated Use {@link #using(McpTransport)} instead. */ + @Deprecated public static McpSyncClient sync(McpTransport transport) { - return new McpSyncClient(async(transport)); + return using(transport).sync(); } /** - * Creates a synchronous MCP client with custom configuration. This method wraps an - * asynchronous client to provide synchronous operations. - * @param transport The transport layer implementation for MCP communication - * @param requestTimeout The duration to wait before timing out requests - * @param objectMapper The Jackson object mapper for JSON - * serialization/deserialization - * @return A new instance of {@link McpSyncClient} + * @deprecated Use {@link #using(McpTransport)} instead. */ - public static McpSyncClient sync(McpTransport transport, Duration requestTimeout, ObjectMapper objectMapper) { - return new McpSyncClient(async(transport, requestTimeout, objectMapper)); + @Deprecated + public static McpSyncClient sync(McpTransport transport, Duration requestTimeout) { + return using(transport).withRequestTimeout(requestTimeout).sync(); } } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java index 2454a2e..6d1e964 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java @@ -17,12 +17,16 @@ package org.springframework.ai.mcp.client; import java.time.Duration; +import java.util.concurrent.TimeoutException; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; -import org.springframework.ai.mcp.client.util.Assert; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.GetPromptRequest; import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult; import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult; +import org.springframework.ai.mcp.util.Assert; /** * @author Dariusz Jędrzejczyk @@ -30,6 +34,8 @@ import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult; */ public class McpSyncClient implements AutoCloseable { + private static final Logger logger = LoggerFactory.getLogger(McpSyncClient.class); + // TODO: Consider providing a client config to set this properly // this is currently a concern only because AutoCloseable is used - perhaps it // is not a requirement? @@ -44,7 +50,18 @@ public class McpSyncClient implements AutoCloseable { @Override public void close() { - this.delegate.closeGracefully().block(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS)); + this.delegate.close(); + } + + public boolean closeGracefully() { + try { + this.delegate.closeGracefully().block(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS)); + } + catch (RuntimeException e) { + logger.warn("Client didn't close within timeout of {} ms.", DEFAULT_CLOSE_TIMEOUT_MS, e); + return false; + } + return true; } /** diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/SseServerTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/SseClientTransport.java similarity index 60% rename from spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/SseServerTransport.java rename to spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/SseClientTransport.java index 504f504..7a16c2d 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/SseServerTransport.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/sse/SseClientTransport.java @@ -16,26 +16,27 @@ package org.springframework.ai.mcp.client.sse; import java.io.IOException; -import java.util.concurrent.Executors; import java.util.function.BiConsumer; -import java.util.function.Consumer; +import java.util.function.Function; +import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.Disposable; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.publisher.Sinks; import reactor.core.publisher.SynchronousSink; -import reactor.core.scheduler.Scheduler; import reactor.core.scheduler.Schedulers; import reactor.util.retry.Retry; import reactor.util.retry.Retry.RetrySignal; -import org.springframework.ai.mcp.client.util.Assert; -import org.springframework.ai.mcp.spec.AbstractMcpTransport; +import org.springframework.ai.mcp.spec.McpError; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; +import org.springframework.ai.mcp.spec.McpTransport; +import org.springframework.ai.mcp.util.Assert; import org.springframework.core.ParameterizedTypeReference; import org.springframework.http.MediaType; import org.springframework.http.codec.ServerSentEvent; @@ -70,9 +71,9 @@ import org.springframework.web.reactive.function.client.WebClient; * "https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse">MCP * HTTP with SSE Transport Specification */ -public class SseServerTransport extends AbstractMcpTransport { +public class SseClientTransport implements McpTransport { - private final static Logger logger = LoggerFactory.getLogger(SseServerTransport.class); + private final static Logger logger = LoggerFactory.getLogger(SseClientTransport.class); /** * Event type for JSON-RPC messages received through the SSE connection. The server @@ -109,13 +110,7 @@ public class SseServerTransport extends AbstractMcpTransport { * ObjectMapper for serializing outbound messages and deserializing inbound messages. * Handles conversion between JSON-RPC messages and their string representation. */ - private ObjectMapper objectMapper; - - /** - * Dedicated scheduler for processing outbound messages. Ensures sequential processing - * of outbound messages on a single thread. - */ - private Scheduler outboundScheduler; + protected ObjectMapper objectMapper; /** * Subscription for the SSE connection handling inbound messages. Used for cleanup @@ -123,12 +118,6 @@ public class SseServerTransport extends AbstractMcpTransport { */ private Disposable inboundSubscription; - /** - * Subscription for the outbound message processing pipeline. Used for cleanup during - * transport shutdown. - */ - private Disposable outboundSubscription; - /** * Flag indicating if the transport is in the process of shutting down. Used to * prevent new operations during shutdown and handle cleanup gracefully. @@ -139,38 +128,100 @@ public class SseServerTransport extends AbstractMcpTransport { * Sink for managing the message endpoint URI provided by the server. Stores the most * recent endpoint URI and makes it available for outbound message processing. */ - protected final Sinks.Many messageEndpointSink = Sinks.many().replay().latest(); + protected final Sinks.One messageEndpointSink = Sinks.one(); /** - * Constructs a new SseServerTransport with the specified WebClient builder. Uses a + * Constructs a new SseClientTransport with the specified WebClient builder. Uses a * default ObjectMapper instance for JSON processing. * @param webClientBuilder the WebClient.Builder to use for creating the WebClient * instance * @throws IllegalArgumentException if webClientBuilder is null */ - public SseServerTransport(WebClient.Builder webClientBuilder) { + public SseClientTransport(WebClient.Builder webClientBuilder) { this(webClientBuilder, new ObjectMapper()); } /** - * Constructs a new SseServerTransport with the specified WebClient builder and + * Constructs a new SseClientTransport with the specified WebClient builder and * ObjectMapper. Initializes both inbound and outbound message processing pipelines. * @param webClientBuilder the WebClient.Builder to use for creating the WebClient * instance * @param objectMapper the ObjectMapper to use for JSON processing * @throws IllegalArgumentException if either parameter is null */ - public SseServerTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) { + public SseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); Assert.notNull(webClientBuilder, "WebClient.Builder must not be null"); this.objectMapper = objectMapper; this.webClient = webClientBuilder.build(); + } - this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "outbound"); + @Override + public Mono connect(Function, Mono> handler) { + Flux> events = eventStream(); + this.inboundSubscription = events.concatMap(event -> Mono.just(event).handle((e, s) -> { + if (ENDPOINT_EVENT_TYPE.equals(event.event())) { + String messageEndpointUri = event.data(); + if (messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) { + s.complete(); + } + else { + // TODO: clarify with the spec if multiple events can be + // received + s.error(new McpError("Failed to handle SSE endpoint event")); + } + } + else if (MESSAGE_EVENT_TYPE.equals(event.event())) { + try { + JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, event.data()); + s.next(message); + } + catch (IOException ioException) { + s.error(ioException); + } + } + else { + s.error(new McpError("Received unrecognized SSE event type: " + event.event())); + } + }).transform(handler)).subscribe(); - startInboundProcessing(); - startOutboundProcessing(); + // The connection is established once the server sends the endpoint event + return messageEndpointSink.asMono().then(); + } + + @Override + public Mono sendMessage(JSONRPCMessage message) { + // The messageEndpoint is the endpoint URI to send the messages + // It is provided by the server as part of the endpoint event + return messageEndpointSink.asMono().flatMap(messageEndpointUri -> { + if (isClosing) { + return Mono.empty(); + } + try { + String jsonText = this.objectMapper.writeValueAsString(message); + return webClient.post() + .uri(messageEndpointUri) + .contentType(MediaType.APPLICATION_JSON) + .bodyValue(jsonText) + .retrieve() + .toBodilessEntity() + .doOnSuccess(response -> { + logger.debug("Message sent successfully"); + }) + .doOnError(error -> { + if (!isClosing) { + logger.error("Error sending message: {}", error.getMessage()); + } + }); + } + catch (IOException e) { + if (!isClosing) { + return Mono.error(new RuntimeException("Failed to serialize message", e)); + } + return Mono.empty(); + } + }).then(); // TODO: Consider non-200-ok response } /** @@ -178,18 +229,15 @@ public class SseServerTransport extends AbstractMcpTransport { * connection and sets up event handling for both message and endpoint events. * Includes automatic retry logic for handling transient connection failures. */ - private void startInboundProcessing() {// @formatter:off - this.inboundSubscription = this.webClient + // visible for tests + protected Flux> eventStream() {// @formatter:off + return this.webClient .get() .uri(SSE_ENDPOINT) .accept(MediaType.TEXT_EVENT_STREAM) .retrieve() .bodyToFlux(SSE_TYPE) - .retryWhen(Retry.from(retrySignal -> retrySignal.handle(inboundRetryHandler))) - .subscribe( - onInboudEvent, - onInboundError, - onInboundComplete); + .retryWhen(Retry.from(retrySignal -> retrySignal.handle(inboundRetryHandler))); } // @formatter:on /** @@ -211,112 +259,6 @@ public class SseServerTransport extends AbstractMcpTransport { sink.error(retrySpec.failure()); }; - /** - * Processes incoming SSE events, handling both message and endpoint events. For - * message events, deserializes and emits JSON-RPC messages to the inbound sink. For - * endpoint events, updates the message endpoint URI for outbound messages. - */ - private Consumer> onInboudEvent = event -> { - if (!isClosing) { - try { - logger.debug("Received SSE event: {}", event); - - // When a client connects, the server MUST send an endpoint event - // containing a URI for the client to use for sending messages. - // All subsequent client messages MUST be sent as HTTP POST requests - // to this endpoint. - if (ENDPOINT_EVENT_TYPE.equals(event.event())) { - - String messageEndpointUri = event.data(); - messageEndpointSink.tryEmitNext(messageEndpointUri); - - } - else if (MESSAGE_EVENT_TYPE.equals(event.event())) { - JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, event.data()); - - if (!this.getInboundSink().tryEmitNext(message).isSuccess()) { - throw new RuntimeException("Failed to enqueue message"); - } - } - } - catch (Exception e) { - if (!isClosing) { - throw new RuntimeException(e); - } - } - } - }; - - /** - * Error handler for the inbound SSE stream. - */ - private Consumer onInboundError = error -> { - if (!isClosing) { - logger.error("Error receiving SSE: {}", error); - } - }; - - /** - * Completion handler for the inbound SSE stream. - */ - private Runnable onInboundComplete = () -> { - if (!isClosing) { - logger.info("SSE stream completed!"); - } - }; - - /** - * Initializes and starts the outbound message processing. Sets up the pipeline for - * sending messages to the server-provided endpoint. Messages are sent as HTTP POST - * requests with JSON content. - * - *

- * - * This pipeline depends on the inbound pipeline to provide the message endpoint URI - * for the client to use for sending messages. The messageEndpointSink is used to - * store the most recent message endpoint URI provided by the server. - */ - private void startOutboundProcessing() {// @formatter:off - this.outboundSubscription = messageEndpointSink.asFlux() - // The messageEndpoint is the endpoint URI to send the messages - // It is provided by the server as part of the endpoint event - .flatMap(messageEndpointUri -> this.getOutboundSink() - .asFlux() - .publishOn(outboundScheduler) - .flatMap(jsonRpcMessage -> { - if (isClosing) { - return Mono.empty(); - } - - try { - String jsonText = this.objectMapper.writeValueAsString(jsonRpcMessage); - - return webClient.post() - .uri(messageEndpointUri) - .contentType(MediaType.APPLICATION_JSON) - .bodyValue(jsonText) - .retrieve() - .toBodilessEntity() - .doOnSuccess(response -> { - logger.debug("Message sent successfully"); - }) - .doOnError(error -> { - if (!isClosing) { - logger.error("Error sending message: {}", error.getMessage()); - } - }) - .thenReturn(jsonRpcMessage); - } - catch (IOException e) { - if (!isClosing) { - return Mono.error(new RuntimeException("Failed to serialize message", e)); - } - return Mono.empty(); - } - })) - .subscribe(); - } // @formatter:on - /** * Implements graceful shutdown of the transport. Cleans up all resources including * subscriptions and schedulers. Ensures orderly shutdown of both inbound and outbound @@ -333,21 +275,15 @@ public class SseServerTransport extends AbstractMcpTransport { if (inboundSubscription != null) { inboundSubscription.dispose(); } - - if (outboundSubscription != null) { - outboundSubscription.dispose(); - } - - // Dispose of scheduler - if (outboundScheduler != null) { - outboundScheduler.dispose(); - } - // Complete the endpoint sink - // messageEndpointSink.tryEmitComplete(); }) .then() .subscribeOn(Schedulers.boundedElastic()); } // @formatter:on + @Override + public T unmarshalFrom(Object data, TypeReference typeRef) { + return this.objectMapper.convertValue(data, typeRef); + } + } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParameters.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParameters.java index 8c01c86..d2b1702 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParameters.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParameters.java @@ -26,7 +26,7 @@ import java.util.stream.Collectors; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; -import org.springframework.ai.mcp.client.util.Assert; +import org.springframework.ai.mcp.util.Assert; /** * Server parameters for stdio client. diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioClientTransport.java similarity index 64% rename from spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java rename to spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioClientTransport.java index 394bdec..1bf5a58 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioClientTransport.java @@ -24,18 +24,23 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; +import java.util.function.Consumer; +import java.util.function.Function; +import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; import reactor.core.scheduler.Scheduler; import reactor.core.scheduler.Schedulers; -import org.springframework.ai.mcp.client.util.Assert; -import org.springframework.ai.mcp.spec.AbstractMcpTransport; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; +import org.springframework.ai.mcp.util.Assert; +import org.springframework.ai.mcp.spec.McpTransport; /** * Implementation of the MCP Stdio transport that communicates with a server process using @@ -45,13 +50,19 @@ import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; * @author Christian Tzolov * @author Dariusz Jędrzejczyk */ -public class StdioServerTransport extends AbstractMcpTransport { +public class StdioClientTransport implements McpTransport { - private static final Logger logger = LoggerFactory.getLogger(StdioServerTransport.class); + private static final Logger logger = LoggerFactory.getLogger(StdioClientTransport.class); + + private final Sinks.Many inboundSink; + + private final Sinks.Many outboundSink; /** The server process being communicated with */ private Process process; + private ObjectMapper objectMapper; + /** Scheduler for handling inbound messages from the server process */ private Scheduler inboundScheduler; @@ -64,28 +75,38 @@ public class StdioServerTransport extends AbstractMcpTransport { /** Parameters for configuring and starting the server process */ private final ServerParameters params; + private final Sinks.Many errorSink; + + // visible for tests + Consumer errorHandler = error -> logger.error("Error received: {}", error); + /** - * Creates a new StdioServerTransport with the specified parameters and default + * Creates a new StdioClientTransport with the specified parameters and default * ObjectMapper. * @param params The parameters for configuring the server process */ - public StdioServerTransport(ServerParameters params) { + public StdioClientTransport(ServerParameters params) { this(params, new ObjectMapper()); } /** - * Creates a new StdioServerTransport with the specified parameters and ObjectMapper. + * Creates a new StdioClientTransport with the specified parameters and ObjectMapper. * @param params The parameters for configuring the server process * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization */ - public StdioServerTransport(ServerParameters params, ObjectMapper objectMapper) { - super(objectMapper); - + public StdioClientTransport(ServerParameters params, ObjectMapper objectMapper) { Assert.notNull(params, "The params can not be null"); Assert.notNull(objectMapper, "The ObjectMapper can not be null"); + this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); + this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); + this.params = params; + this.objectMapper = objectMapper; + + this.errorSink = Sinks.many().unicast().onBackpressureBuffer(); + // Start threads this.inboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "inbound"); this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "outbound"); @@ -100,39 +121,39 @@ public class StdioServerTransport extends AbstractMcpTransport { * are null */ @Override - public void start() { - // Let's kick off the abstraction layer that will consume the logical messages - // pushed via sinks. The code that follows is actually feeding the sinks with - // data. - super.start(); + public Mono connect(Function, Mono> handler) { + return Mono.fromRunnable(() -> { + handleIncomingMessages(handler); + handleIncomingErrors(); - // Prepare command and environment - List fullCommand = new ArrayList<>(); - fullCommand.add(params.getCommand()); - fullCommand.addAll(params.getArgs()); + // Prepare command and environment + List fullCommand = new ArrayList<>(); + fullCommand.add(params.getCommand()); + fullCommand.addAll(params.getArgs()); - ProcessBuilder processBuilder = this.getProcessBuilder(); - processBuilder.command(fullCommand); - processBuilder.environment().putAll(params.getEnv()); + ProcessBuilder processBuilder = this.getProcessBuilder(); + processBuilder.command(fullCommand); + processBuilder.environment().putAll(params.getEnv()); - // Start the process - try { - this.process = processBuilder.start(); - } - catch (IOException e) { - throw new RuntimeException("Failed to start process with command: " + fullCommand, e); - } + // Start the process + try { + this.process = processBuilder.start(); + } + catch (IOException e) { + throw new RuntimeException("Failed to start process with command: " + fullCommand, e); + } - // Validate process streams - if (this.process.getInputStream() == null || process.getOutputStream() == null) { - this.process.destroy(); - throw new RuntimeException("Process input or output stream is null"); - } + // Validate process streams + if (this.process.getInputStream() == null || process.getOutputStream() == null) { + this.process.destroy(); + throw new RuntimeException("Process input or output stream is null"); + } - // Start threads - startInboundProcessing(); - startOutboundProcessing(); - startErrorProcessing(); + // Start threads + startInboundProcessing(); + startOutboundProcessing(); + startErrorProcessing(); + }).subscribeOn(Schedulers.boundedElastic()); } /** @@ -144,6 +165,19 @@ public class StdioServerTransport extends AbstractMcpTransport { return new ProcessBuilder(); } + /** + * Sets the handler for processing transport-level errors. + * + *

+ * The provided handler will be called when errors occur during transport operations, + * such as connection failures or protocol violations. + *

+ * @param errorHandler a consumer that processes error messages + */ + public void setInboundErrorHandler(Consumer errorHandler) { + this.errorHandler = errorHandler; + } + /** * Waits for the server process to exit. * @throws RuntimeException if the process is interrupted while waiting @@ -170,7 +204,7 @@ public class StdioServerTransport extends AbstractMcpTransport { try { logger.error("Received error line: {}", line); // TODO: handle errors, etc. - this.getErrorSink().tryEmitNext(line); + this.errorSink.tryEmitNext(line); } catch (Exception e) { throw new RuntimeException(e); @@ -183,6 +217,35 @@ public class StdioServerTransport extends AbstractMcpTransport { }); } + private void handleIncomingMessages(Function, Mono> inboundMessageHandler) { + this.inboundSink.asFlux() + .flatMap(message -> Mono.just(message) + .transform(inboundMessageHandler) + .contextWrite(ctx -> ctx.put("observation", "myObservation"))) + .subscribe(); + } + + private void handleIncomingErrors() { + this.errorSink.asFlux().subscribe(e -> { + this.errorHandler.accept(e); + }); + } + + @Override + public Mono sendMessage(JSONRPCMessage message) { + if (this.outboundSink.tryEmitNext(message).isSuccess()) { + // TODO: essentially we could reschedule ourselves in some time and make + // another attempt with the already read data but pause reading until + // success + // In this approach we delegate the retry and the backpressure onto the + // caller. This might be enough for most cases. + return Mono.empty(); + } + else { + return Mono.error(new RuntimeException("Failed to enqueue message")); + } + } + /** * Starts the inbound processing thread that reads JSON-RPC messages from the * process's input stream. Messages are deserialized and emitted to the inbound sink. @@ -194,7 +257,7 @@ public class StdioServerTransport extends AbstractMcpTransport { while ((line = processReader.readLine()) != null) { try { JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line); - if (!this.getInboundSink().tryEmitNext(message).isSuccess()) { + if (!this.inboundSink.tryEmitNext(message).isSuccess()) { // TODO: Back off, reschedule, give up? throw new RuntimeException("Failed to enqueue message"); } @@ -216,8 +279,7 @@ public class StdioServerTransport extends AbstractMcpTransport { * delimiter. */ private void startOutboundProcessing() { - this.getOutboundSink() - .asFlux() + this.handleOutbound(messages -> messages // this bit is important since writes come from user threads and we // want to ensure that the actual writing happens on a dedicated thread .publishOn(outboundScheduler) @@ -234,8 +296,11 @@ public class StdioServerTransport extends AbstractMcpTransport { s.error(new RuntimeException(e)); } } - }) - .subscribe(); + })); + } + + protected void handleOutbound(Function, Flux> outboundConsumer) { + outboundConsumer.apply(outboundSink.asFlux()).subscribe(); } /** @@ -268,4 +333,13 @@ public class StdioServerTransport extends AbstractMcpTransport { })).then().subscribeOn(Schedulers.boundedElastic()); } + public Sinks.Many getErrorSink() { + return this.errorSink; + } + + @Override + public T unmarshalFrom(Object data, TypeReference typeRef) { + return this.objectMapper.convertValue(data, typeRef); + } + } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java deleted file mode 100644 index 0f0556e..0000000 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java +++ /dev/null @@ -1,122 +0,0 @@ -/* - * Copyright 2024-2024 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.ai.mcp.spec; - -import java.util.function.Consumer; - -import com.fasterxml.jackson.databind.ObjectMapper; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import reactor.core.publisher.Mono; -import reactor.core.publisher.Sinks; - -import org.springframework.ai.mcp.client.util.Assert; -import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; - -/** - * @author Christian Tzolov - * @author Dariusz Jędrzejczyk - */ -// TODO: The API is not finalized. Implementing SSE will dictate a more concrete -// structure and relationship between the base and the actual implementation including -// proper encapsulation and state management. -public abstract class AbstractMcpTransport implements McpTransport { - - private static final Logger logger = LoggerFactory.getLogger(AbstractMcpTransport.class); - - protected final ObjectMapper objectMapper; - - private final Sinks.Many errorSink; - - private final Sinks.Many inboundSink; - - private final Sinks.Many outboundSink; - - private Consumer inboundMessageHandler = message -> logger.debug("Received message: {}", message); - - private Consumer errorHandler = error -> logger.error("Received error: {}", error); - - public AbstractMcpTransport() { - this(new ObjectMapper()); - } - - public AbstractMcpTransport(ObjectMapper objectMapper) { - - Assert.notNull(objectMapper, "ObjectMapper must not be null"); - this.objectMapper = objectMapper; - - this.errorSink = Sinks.many().unicast().onBackpressureBuffer(); - this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); - this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); - } - - public ObjectMapper getObjectMapper() { - return this.objectMapper; - } - - private void handleIncomingMessages() { - this.inboundSink.asFlux().subscribe(message -> this.inboundMessageHandler.accept(message)); - } - - private void handleIncomingErrors() { - this.errorSink.asFlux().subscribe(e -> { - this.errorHandler.accept(e); - }); - } - - protected Sinks.Many getInboundSink() { - return inboundSink; - } - - protected Sinks.Many getOutboundSink() { - return outboundSink; - } - - protected Sinks.Many getErrorSink() { - return errorSink; - } - - public void setInboundMessageHandler(Consumer inboundMessageHandler) { - this.inboundMessageHandler = inboundMessageHandler; - } - - public void setInboundErrorHandler(Consumer errorHandler) { - this.errorHandler = errorHandler; - } - - // Start processing incoming messages - public void start() { - this.handleIncomingErrors(); - this.handleIncomingMessages(); - } - - @Override - public Mono sendMessage(JSONRPCMessage message) { - if (this.outboundSink.tryEmitNext(message).isSuccess()) { - // TODO: essentially we could reschedule ourselves in some time and make - // another attempt with the already read data but pause reading until - // success - // In this approach we delegate the retry and the backpressure onto the - // caller. This might be enough for most cases. - return Mono.empty(); - } - else { - return Mono.error(new RuntimeException("Failed to enqueue message")); - } - } - -} diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java index fdbadc0..cb82f54 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java @@ -23,13 +23,13 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.Disposable; import reactor.core.publisher.Mono; import reactor.core.publisher.MonoSink; -import org.springframework.ai.mcp.client.util.Assert; +import org.springframework.ai.mcp.util.Assert; /** * Default implementation of the MCP (Model Context Protocol) session that manages @@ -56,9 +56,6 @@ public class DefaultMcpSession implements McpSession { /** Duration to wait for request responses before timing out */ private final Duration requestTimeout; - /** JSON object mapper for serialization/deserialization */ - private final ObjectMapper objectMapper; - /** Transport layer implementation for message exchange */ private final McpTransport transport; @@ -77,6 +74,8 @@ public class DefaultMcpSession implements McpSession { /** Atomic counter for generating unique request IDs */ private final AtomicLong requestCounter = new AtomicLong(0); + private final Disposable connection; + /** * Functional interface for handling incoming JSON-RPC requests. Implementations * should process the request parameters and return a response. @@ -112,37 +111,36 @@ public class DefaultMcpSession implements McpSession { /** * Creates a new DefaultMcpSession with the specified configuration. * @param requestTimeout Duration to wait for responses - * @param objectMapper JSON object mapper for message serialization * @param transport Transport implementation for message exchange */ - public DefaultMcpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport) { - this(requestTimeout, objectMapper, transport, Map.of(), Map.of()); + public DefaultMcpSession(Duration requestTimeout, McpTransport transport) { + this(requestTimeout, transport, Map.of(), Map.of()); } /** * Creates a new DefaultMcpSession with the specified configuration and handlers. * @param requestTimeout Duration to wait for responses - * @param objectMapper JSON object mapper for message serialization * @param transport Transport implementation for message exchange * @param requestHandlers Map of method names to request handlers * @param notificationHandlers Map of method names to notification handlers */ - public DefaultMcpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport, + public DefaultMcpSession(Duration requestTimeout, McpTransport transport, Map requestHandlers, Map notificationHandlers) { - Assert.notNull(objectMapper, "The ObjectMapper can not be null"); Assert.notNull(requestTimeout, "The requstTimeout can not be null"); Assert.notNull(transport, "The transport can not be null"); Assert.notNull(requestHandlers, "The requestHandlers can not be null"); Assert.notNull(notificationHandlers, "The notificationHandlers can not be null"); this.requestTimeout = requestTimeout; - this.objectMapper = objectMapper; this.transport = transport; this.requestHandlers.putAll(requestHandlers); this.notificationHandlers.putAll(notificationHandlers); - this.transport.setInboundMessageHandler(message -> { + // TODO: consider mono.transformDeferredContextual where the Context contains the + // Observation associated with the individual message - it can be used to + // create child Observation and emit it together with the message to the consumer + this.connection = this.transport.connect(mono -> mono.doOnNext(message -> { if (message instanceof McpSchema.JSONRPCResponse response) { var sink = pendingResponses.remove(response.id()); if (sink == null) { @@ -165,11 +163,7 @@ public class DefaultMcpSession implements McpSession { handleIncomingNotification(notification).subscribe(null, error -> logger.error("Error handling notification: {}", error.getMessage())); } - }); - - this.transport.setInboundErrorHandler(error -> logger.error("Error received: {}", error)); - - this.transport.start(); + })).subscribe(); } /** @@ -252,7 +246,7 @@ public class DefaultMcpSession implements McpSession { sink.complete(); } else { - sink.next(this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef)); + sink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef)); } } }); @@ -277,6 +271,7 @@ public class DefaultMcpSession implements McpSession { */ @Override public Mono closeGracefully() { + this.connection.dispose(); return transport.closeGracefully(); } @@ -285,6 +280,7 @@ public class DefaultMcpSession implements McpSession { */ @Override public void close() { + this.connection.dispose(); transport.close(); } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java index 5fc6014..f81a6c6 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java @@ -16,8 +16,9 @@ package org.springframework.ai.mcp.spec; -import java.util.function.Consumer; +import java.util.function.Function; +import com.fasterxml.jackson.core.type.TypeReference; import reactor.core.publisher.Mono; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; @@ -59,7 +60,7 @@ public interface McpTransport { * necessary resources and establishes the connection to the server. *

*/ - void start(); + Mono connect(Function, Mono> handler); /** * Closes the transport connection and releases any associated resources. @@ -80,30 +81,6 @@ public interface McpTransport { */ Mono closeGracefully(); - /** - * Sets the handler for processing incoming messages from the server. - * - *

- * The provided handler will be called whenever a new message is received from the - * server. Messages are expected to follow the JSON-RPC format as defined in the MCP - * specification. - *

- * @param inboundMessageHandler a consumer that processes incoming - * {@link JSONRPCMessage}s - */ - void setInboundMessageHandler(Consumer inboundMessageHandler); - - /** - * Sets the handler for processing transport-level errors. - * - *

- * The provided handler will be called when errors occur during transport operations, - * such as connection failures or protocol violations. - *

- * @param inboundErrorHandler a consumer that processes error messages - */ - void setInboundErrorHandler(Consumer inboundErrorHandler); - /** * Sends a message to the server asynchronously. * @@ -116,4 +93,6 @@ public interface McpTransport { */ Mono sendMessage(JSONRPCMessage message); + T unmarshalFrom(Object data, TypeReference typeRef); + } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/util/Assert.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/util/Assert.java similarity index 98% rename from spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/util/Assert.java rename to spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/util/Assert.java index efcc26a..256d063 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/util/Assert.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/util/Assert.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.mcp.client.util; +package org.springframework.ai.mcp.util; import reactor.util.annotation.Nullable; diff --git a/spring-ai-mcp-core/src/main/resources/logback.xml b/spring-ai-mcp-core/src/main/resources/logback.xml index 8cea17b..84e350d 100644 --- a/spring-ai-mcp-core/src/main/resources/logback.xml +++ b/spring-ai-mcp-core/src/main/resources/logback.xml @@ -20,7 +20,7 @@ - + diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests2.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests2.java index 57a4ba3..12c3d71 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests2.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests2.java @@ -18,12 +18,10 @@ package org.springframework.ai.mcp.attic; import java.time.Duration; -import com.fasterxml.jackson.databind.ObjectMapper; - import org.springframework.ai.mcp.client.McpClient; import org.springframework.ai.mcp.client.McpSyncClient; import org.springframework.ai.mcp.client.stdio.ServerParameters; -import org.springframework.ai.mcp.client.stdio.StdioServerTransport; +import org.springframework.ai.mcp.client.stdio.StdioClientTransport; import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult; import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult; import org.springframework.ai.mcp.spec.McpSchema.Resource; @@ -40,8 +38,9 @@ public class ClientSessionTests2 { .args("--directory", "dir", "run", "mcp-server-sqlite", "--db-path", "~/test.db") .build(); - try (McpSyncClient clientSession = McpClient.sync(new StdioServerTransport(stdioParams), Duration.ofSeconds(10), - new ObjectMapper())) { + try (McpSyncClient clientSession = McpClient.using(new StdioClientTransport(stdioParams)) + .withRequestTimeout(Duration.ofSeconds(10)) + .sync()) { clientSession.initialize(); diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests3.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests3.java index e4bd3e1..de1417c 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests3.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ClientSessionTests3.java @@ -18,12 +18,10 @@ package org.springframework.ai.mcp.attic; import java.time.Duration; -import com.fasterxml.jackson.databind.ObjectMapper; - import org.springframework.ai.mcp.client.McpClient; import org.springframework.ai.mcp.client.McpSyncClient; import org.springframework.ai.mcp.client.stdio.ServerParameters; -import org.springframework.ai.mcp.client.stdio.StdioServerTransport; +import org.springframework.ai.mcp.client.stdio.StdioClientTransport; import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult; import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult; import org.springframework.ai.mcp.spec.McpSchema.Resource; @@ -43,8 +41,9 @@ public class ClientSessionTests3 { McpSyncClient clientSession = null; try { - clientSession = McpClient.sync(new StdioServerTransport(stdioParams), Duration.ofSeconds(10), - new ObjectMapper()); + clientSession = McpClient.using(new StdioClientTransport(stdioParams)) + .withRequestTimeout(Duration.ofSeconds(10)) + .sync(); clientSession.initialize(); diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/MainSSE.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/MainSSE.java index 735cf7b..1e96dad 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/MainSSE.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/MainSSE.java @@ -18,11 +18,9 @@ package org.springframework.ai.mcp.attic; import java.time.Duration; import java.util.Map; -import com.fasterxml.jackson.databind.ObjectMapper; - import org.springframework.ai.mcp.client.McpClient; import org.springframework.ai.mcp.client.McpSyncClient; -import org.springframework.ai.mcp.client.sse.SseServerTransport; +import org.springframework.ai.mcp.client.sse.SseClientTransport; import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest; import org.springframework.ai.mcp.spec.McpSchema.CallToolResult; import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult; @@ -40,9 +38,10 @@ public class MainSSE { public static void main(String[] args) { - try (McpSyncClient mcpClient = McpClient.sync( - new SseServerTransport(WebClient.builder().baseUrl("http://localhost:3001")), Duration.ofSeconds(1000), - new ObjectMapper())) { + try (McpSyncClient mcpClient = McpClient + .using(new SseClientTransport(WebClient.builder().baseUrl("http://localhost:3001"))) + .withRequestTimeout(Duration.ofSeconds(1000)) + .sync()) { mcpClient.initialize(); diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParametersParser.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ServerParametersParser.java similarity index 94% rename from spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParametersParser.java rename to spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ServerParametersParser.java index 723dde2..e093723 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/ServerParametersParser.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/attic/ServerParametersParser.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.mcp.client.stdio; +package org.springframework.ai.mcp.attic; import java.util.List; import java.util.Map; @@ -23,7 +23,7 @@ import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.databind.ObjectMapper; -import org.springframework.ai.mcp.client.stdio.ServerParametersParser.McpServerConfigurations.McpServerConfiguration; +import org.springframework.ai.mcp.attic.ServerParametersParser.McpServerConfigurations.McpServerConfiguration; /** * @author Christian Tzolov diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java index 07adc1c..08cf29a 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java @@ -19,7 +19,6 @@ package org.springframework.ai.mcp.client; import java.time.Duration; import java.util.Map; -import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -53,7 +52,7 @@ public abstract class AbstractMcpAsyncClientTests { private static final String ECHO_TEST_MESSAGE = "Hello MCP Spring AI!"; - abstract protected void createMcpTransport(); + abstract protected McpTransport createMcpTransport(); abstract protected void onStart(); @@ -62,10 +61,10 @@ public abstract class AbstractMcpAsyncClientTests { @BeforeEach void setUp() { onStart(); - createMcpTransport(); + this.mcpTransport = createMcpTransport(); assertThatCode(() -> { - mcpAsyncClient = McpClient.async(mcpTransport, TIMEOUT, new ObjectMapper()); + mcpAsyncClient = McpClient.using(mcpTransport).withRequestTimeout(TIMEOUT).async(); mcpAsyncClient.initialize().block(Duration.ofSeconds(10)); }).doesNotThrowAnyException(); } @@ -81,14 +80,12 @@ public abstract class AbstractMcpAsyncClientTests { @Test void testConstructorWithInvalidArguments() { - assertThatThrownBy(() -> McpClient.sync(null, TIMEOUT, new ObjectMapper())) - .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> McpClient.using(null).sync()).isInstanceOf(IllegalArgumentException.class) + .hasMessage("Transport must not be null"); - assertThatThrownBy(() -> McpClient.sync(mcpTransport, null, new ObjectMapper())) - .isInstanceOf(IllegalArgumentException.class); - - assertThatThrownBy(() -> McpClient.sync(mcpTransport, TIMEOUT, null)) - .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> McpClient.using(mcpTransport).withRequestTimeout(null).sync()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Request timeout must not be null"); } @Test diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java index a66bc99..00dc68d 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java @@ -19,7 +19,6 @@ package org.springframework.ai.mcp.client; import java.time.Duration; import java.util.Map; -import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -53,7 +52,7 @@ public abstract class AbstractMcpSyncClientTests { protected McpTransport mcpTransport; - abstract protected void createMcpTransport(); + abstract protected McpTransport createMcpTransport(); abstract protected void onStart(); @@ -62,10 +61,10 @@ public abstract class AbstractMcpSyncClientTests { @BeforeEach void setUp() { onStart(); - createMcpTransport(); + this.mcpTransport = createMcpTransport(); assertThatCode(() -> { - mcpSyncClient = McpClient.sync(mcpTransport, TIMEOUT, new ObjectMapper()); + mcpSyncClient = McpClient.using(mcpTransport).withRequestTimeout(TIMEOUT).sync(); mcpSyncClient.initialize(); }).doesNotThrowAnyException(); } @@ -80,14 +79,12 @@ public abstract class AbstractMcpSyncClientTests { @Test void testConstructorWithInvalidArguments() { - assertThatThrownBy(() -> McpClient.sync(null, TIMEOUT, new ObjectMapper())) - .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> McpClient.using(null).sync()).isInstanceOf(IllegalArgumentException.class) + .hasMessage("Transport must not be null"); - assertThatThrownBy(() -> McpClient.sync(mcpTransport, null, new ObjectMapper())) - .isInstanceOf(IllegalArgumentException.class); - - assertThatThrownBy(() -> McpClient.sync(mcpTransport, TIMEOUT, null)) - .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> McpClient.using(mcpTransport).withRequestTimeout(null).sync()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Request timeout must not be null"); } @Test diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseServerTransportTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseClientTransportTests.java similarity index 85% rename from spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseServerTransportTests.java rename to spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseClientTransportTests.java index 5955c2d..6e2e378 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseServerTransportTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseClientTransportTests.java @@ -19,21 +19,23 @@ package org.springframework.ai.mcp.client.sse; import java.time.Duration; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; import com.fasterxml.jackson.databind.ObjectMapper; -import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.testcontainers.containers.GenericContainer; import org.testcontainers.containers.wait.strategy.Wait; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; import reactor.test.StepVerifier; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest; +import org.springframework.http.codec.ServerSentEvent; import org.springframework.web.reactive.function.client.WebClient; import static org.assertj.core.api.Assertions.assertThat; @@ -41,12 +43,12 @@ import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.Assertions.assertThatThrownBy; /** - * Tests for the {@link SseServerTransport} class. + * Tests for the {@link SseClientTransport} class. * * @author Christian Tzolov */ @Timeout(15) -class SseServerTransportTests { +class SseClientTransportTests { static String host = "http://localhost:3001"; @@ -63,37 +65,44 @@ class SseServerTransportTests { private ObjectMapper objectMapper; // Test class to access protected methods - static class TestSseServerTransport extends SseServerTransport { + static class TestSseServerTransport extends SseClientTransport { private final AtomicInteger inboundMessageCount = new AtomicInteger(0); + private Sinks.Many> events = Sinks.many().unicast().onBackpressureBuffer(); + public TestSseServerTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) { super(webClientBuilder, objectMapper); } - // simulate receiving a message from the server - public void emitTestMessage(McpSchema.JSONRPCMessage message) { - getInboundSink().tryEmitNext(message); - inboundMessageCount.incrementAndGet(); + // @Override + // public Mono connect(Function, + // Mono> handler) { + // simulateEndpointEvent("https://localhost:3001"); + // return super.connect(handler); + // } + + @Override + protected Flux> eventStream() { + return super.eventStream().mergeWith(events.asFlux()); } public String getLastEndpoint() { - return messageEndpointSink.asFlux().blockFirst(); + return messageEndpointSink.asMono().block(); } public int getInboundMessageCount() { return inboundMessageCount.get(); } + public void simulateEndpointEvent(String jsonMessage) { + events.tryEmitNext(ServerSentEvent.builder().event("endpoint").data(jsonMessage).build()); + inboundMessageCount.incrementAndGet(); + } + public void simulateMessageEvent(String jsonMessage) { - try { - McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(getObjectMapper(), jsonMessage); - getInboundSink().tryEmitNext(message); - inboundMessageCount.incrementAndGet(); - } - catch (Exception e) { - throw new RuntimeException("Failed to simulate message event", e); - } + events.tryEmitNext(ServerSentEvent.builder().event("message").data(jsonMessage).build()); + inboundMessageCount.incrementAndGet(); } } @@ -110,6 +119,7 @@ class SseServerTransportTests { webClientBuilder = WebClient.builder().baseUrl(host); objectMapper = new ObjectMapper(); transport = new TestSseServerTransport(webClientBuilder, objectMapper); + transport.connect(Function.identity()).block(); } @AfterEach @@ -131,10 +141,10 @@ class SseServerTransportTests { @Test void constructorValidation() { - assertThatThrownBy(() -> new SseServerTransport(null)).isInstanceOf(IllegalArgumentException.class) + assertThatThrownBy(() -> new SseClientTransport(null)).isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("WebClient.Builder must not be null"); - assertThatThrownBy(() -> new SseServerTransport(webClientBuilder, null)) + assertThatThrownBy(() -> new SseClientTransport(webClientBuilder, null)) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("ObjectMapper must not be null"); } @@ -242,7 +252,7 @@ class SseServerTransportTests { // Create a WebClient that simulates connection failures WebClient.Builder failingWebClientBuilder = WebClient.builder().baseUrl("http://non-existent-host"); - SseServerTransport failingTransport = new SseServerTransport(failingWebClientBuilder); + SseClientTransport failingTransport = new SseClientTransport(failingWebClientBuilder); // Verify that the transport attempts to reconnect StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete(); diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpAsyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpAsyncClientTests.java index 06d91fc..4faa70d 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpAsyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpAsyncClientTests.java @@ -22,10 +22,11 @@ import org.testcontainers.containers.wait.strategy.Wait; import org.springframework.ai.mcp.client.AbstractMcpAsyncClientTests; import org.springframework.ai.mcp.client.McpAsyncClient; +import org.springframework.ai.mcp.spec.McpTransport; import org.springframework.web.reactive.function.client.WebClient; /** - * Tests for the {@link McpAsyncClient} with {@link SseServerTransport}. + * Tests for the {@link McpAsyncClient} with {@link SseClientTransport}. * * @author Christian Tzolov */ @@ -42,8 +43,8 @@ class SseMcpAsyncClientTests extends AbstractMcpAsyncClientTests { .waitingFor(Wait.forHttp("/").forStatusCode(404)); @Override - protected void createMcpTransport() { - this.mcpTransport = new SseServerTransport(WebClient.builder().baseUrl(host)); + protected McpTransport createMcpTransport() { + return new SseClientTransport(WebClient.builder().baseUrl(host)); } @Override diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpSyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpSyncClientTests.java index 823a74e..65e81e0 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpSyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/sse/SseMcpSyncClientTests.java @@ -16,18 +16,17 @@ package org.springframework.ai.mcp.client.sse; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Timeout; import org.testcontainers.containers.GenericContainer; import org.testcontainers.containers.wait.strategy.Wait; import org.springframework.ai.mcp.client.AbstractMcpSyncClientTests; import org.springframework.ai.mcp.client.McpSyncClient; +import org.springframework.ai.mcp.spec.McpTransport; import org.springframework.web.reactive.function.client.WebClient; /** - * Tests for the {@link McpSyncClient} with {@link SseServerTransport}. + * Tests for the {@link McpSyncClient} with {@link SseClientTransport}. * * @author Christian Tzolov */ @@ -44,8 +43,8 @@ class SseMcpSyncClientTests extends AbstractMcpSyncClientTests { .waitingFor(Wait.forHttp("/").forStatusCode(404)); @Override - protected void createMcpTransport() { - this.mcpTransport = new SseServerTransport(WebClient.builder().baseUrl(host)); + protected McpTransport createMcpTransport() { + return new SseClientTransport(WebClient.builder().baseUrl(host)); } @Override diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpAsyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpAsyncClientTests.java index 7e72f8b..64b3307 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpAsyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpAsyncClientTests.java @@ -19,9 +19,10 @@ package org.springframework.ai.mcp.client.stdio; import org.junit.jupiter.api.Timeout; import org.springframework.ai.mcp.client.AbstractMcpAsyncClientTests; +import org.springframework.ai.mcp.spec.McpTransport; /** - * Tests for the {@link McpAsyncClient} with {@link StdioServerTransport}. + * Tests for the {@link McpAsyncClient} with {@link StdioClientTransport}. * * @author Christian Tzolov * @author Dariusz Jędrzejczyk @@ -30,11 +31,11 @@ import org.springframework.ai.mcp.client.AbstractMcpAsyncClientTests; class McpAsyncClientTests extends AbstractMcpAsyncClientTests { @Override - protected void createMcpTransport() { + protected McpTransport createMcpTransport() { ServerParameters stdioParams = ServerParameters.builder("npx") .args("-y", "@modelcontextprotocol/server-everything", "dir") .build(); - this.mcpTransport = new StdioServerTransport(stdioParams); + return new StdioClientTransport(stdioParams); } @Override diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpSyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpSyncClientTests.java index 1a76226..beff095 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpSyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/stdio/McpSyncClientTests.java @@ -16,12 +16,19 @@ package org.springframework.ai.mcp.client.stdio; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.springframework.ai.mcp.client.AbstractMcpSyncClientTests; +import org.springframework.ai.mcp.client.McpSyncClient; +import org.springframework.ai.mcp.spec.McpTransport; + +import static org.assertj.core.api.Assertions.assertThat; /** - * Tests for the {@link McpSyncClient} with {@link StdioServerTransport}. + * Tests for the {@link McpSyncClient} with {@link StdioClientTransport}. * * @author Christian Tzolov * @author Dariusz Jędrzejczyk @@ -30,21 +37,31 @@ import org.springframework.ai.mcp.client.AbstractMcpSyncClientTests; class McpSyncClientTests extends AbstractMcpSyncClientTests { @Override - protected void createMcpTransport() { + protected McpTransport createMcpTransport() { ServerParameters stdioParams = ServerParameters.builder("npx") .args("-y", "@modelcontextprotocol/server-everything", "dir") .build(); - this.mcpTransport = new StdioServerTransport(stdioParams); + + return new StdioClientTransport(stdioParams); + } + + @Test + void customErrorHandlerShouldReceiveErrors() { + AtomicReference receivedError = new AtomicReference<>(); + ((StdioClientTransport) mcpTransport).errorHandler = receivedError::set; + + String errorMessage = "Test error"; + ((StdioClientTransport) mcpTransport).getErrorSink().tryEmitNext(errorMessage); + + assertThat(receivedError.get()).isNotNull().isEqualTo(errorMessage); } @Override protected void onStart() { - } @Override protected void onClose() { - } } diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java index 3428829..75e82a3 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java @@ -19,7 +19,7 @@ package org.springframework.ai.mcp.spec; import java.time.Duration; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; @@ -28,10 +28,12 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; +import reactor.core.scheduler.Schedulers; import reactor.test.StepVerifier; -import org.springframework.ai.mcp.client.McpAsyncClient; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification; import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest; @@ -63,33 +65,49 @@ class DefaultMcpSessionTests { private ObjectMapper objectMapper; @SuppressWarnings("unused") - private static class MockMcpTransport extends AbstractMcpTransport { + private static class MockMcpTransport implements McpTransport { private final AtomicInteger inboundMessageCount = new AtomicInteger(0); - private final AtomicReference lastSentMessage = new AtomicReference<>(); + private Sinks.Many outgoing = Sinks.many().multicast().onBackpressureBuffer(); + + private Sinks.Many inbound = Sinks.many().unicast().onBackpressureBuffer(); + + private Flux outboundView = outgoing.asFlux().cache(1); public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) { + if (inbound.tryEmitNext(message).isFailure()) { + throw new RuntimeException("Failed to emit message " + message); + } inboundMessageCount.incrementAndGet(); - this.getInboundSink().tryEmitNext(message); } @Override public Mono sendMessage(McpSchema.JSONRPCMessage message) { - lastSentMessage.set(message); + if (outgoing.tryEmitNext(message).isFailure()) { + return Mono.error(new RuntimeException("Can't emit outgoing message " + message)); + } return Mono.empty(); } public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() { - return (JSONRPCRequest) lastSentMessage.get(); + return (JSONRPCRequest) outboundView.blockFirst(); } public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() { - return (JSONRPCNotification) lastSentMessage.get(); + return (JSONRPCNotification) outboundView.blockFirst(); } public McpSchema.JSONRPCMessage getLastSentMessage() { - return lastSentMessage.get(); + return outboundView.blockFirst(); + } + + @Override + public Mono connect(Function, Mono> handler) { + return inbound.asFlux() + .publishOn(Schedulers.boundedElastic()) + .flatMap(message -> Mono.just(message).transform(handler)) + .then(); } @Override @@ -97,13 +115,18 @@ class DefaultMcpSessionTests { return Mono.empty(); } + @Override + public T unmarshalFrom(Object data, TypeReference typeRef) { + return new ObjectMapper().convertValue(data, typeRef); + } + } @BeforeEach void setUp() { objectMapper = new ObjectMapper(); transport = new MockMcpTransport(); - session = new DefaultMcpSession(TIMEOUT, objectMapper, transport, Map.of(), + session = new DefaultMcpSession(TIMEOUT, transport, Map.of(), Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params)))); } @@ -116,16 +139,10 @@ class DefaultMcpSessionTests { @Test void testConstructorWithInvalidArguments() { - assertThatThrownBy(() -> new DefaultMcpSession(null, objectMapper, transport)) - .isInstanceOf(IllegalArgumentException.class) + assertThatThrownBy(() -> new DefaultMcpSession(null, transport)).isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("requstTimeout can not be null"); - assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, null, transport)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("ObjectMapper can not be null"); - - assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, objectMapper, null)) - .isInstanceOf(IllegalArgumentException.class) + assertThatThrownBy(() -> new DefaultMcpSession(TIMEOUT, null)).isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("transport can not be null"); } @@ -138,18 +155,13 @@ class DefaultMcpSessionTests { String responseData = "test response"; // Create a Mono that will emit the response after the request is sent - Mono responseMono = session.sendRequest(TEST_METHOD, testParam, responseType) - .doOnSubscribe(subscription -> { - // Wait a bit to ensure the request is sent and captured - Mono.delay(Duration.ofMillis(100)).subscribe(ignored -> { - McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); - transport.simulateIncomingMessage( - new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null)); - }); - }); - + Mono responseMono = session.sendRequest(TEST_METHOD, testParam, responseType); // Verify response handling - StepVerifier.create(responseMono).consumeNextWith(response -> { + StepVerifier.create(responseMono).then(() -> { + McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); + transport.simulateIncomingMessage( + new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null)); + }).consumeNextWith(response -> { // Verify the request was sent McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest(); assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCRequest.class); @@ -162,26 +174,21 @@ class DefaultMcpSessionTests { @Test void testSendRequestWithError() { - Mono responseMono = session.sendRequest(TEST_METHOD, "test", responseType) - .doOnSubscribe(subscription -> { - // Wait a bit to ensure the request is sent and captured - Mono.delay(Duration.ofMillis(100)).subscribe(ignored -> { - McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); - // Simulate error response - McpSchema.JSONRPCResponse.JSONRPCError error = new McpSchema.JSONRPCResponse.JSONRPCError( - McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null); - transport.simulateIncomingMessage( - new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error)); - }); - }); + Mono responseMono = session.sendRequest(TEST_METHOD, "test", responseType); // Verify error handling - StepVerifier.create(responseMono).expectError(McpError.class).verify(); + StepVerifier.create(responseMono).then(() -> { + McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); + // Simulate error response + McpSchema.JSONRPCResponse.JSONRPCError error = new McpSchema.JSONRPCResponse.JSONRPCError( + McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null); + transport.simulateIncomingMessage( + new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error)); + }).expectError(McpError.class).verify(); } @Test void testRequestTimeout() { - Mono responseMono = session.sendRequest(TEST_METHOD, "test", responseType); // Verify timeout @@ -210,7 +217,8 @@ class DefaultMcpSessionTests { String echoMessage = "Hello MCP!"; Map requestHandlers = Map.of(ECHO_METHOD, params -> Mono.just(params)); - session = new DefaultMcpSession(TIMEOUT, objectMapper, transport, requestHandlers, Map.of()); + transport = new MockMcpTransport(); + session = new DefaultMcpSession(TIMEOUT, transport, requestHandlers, Map.of()); // Simulate incoming request McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD, @@ -227,10 +235,11 @@ class DefaultMcpSessionTests { @Test void testNotificationHandling() { - AtomicReference receivedParams = new AtomicReference<>(); + Sinks.One receivedParams = Sinks.one(); - session = new DefaultMcpSession(TIMEOUT, objectMapper, transport, Map.of(), - Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.set(params)))); + transport = new MockMcpTransport(); + session = new DefaultMcpSession(TIMEOUT, transport, Map.of(), + Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params)))); // Simulate incoming notification from the server Map notificationParams = Map.of("status", "ready"); @@ -241,7 +250,7 @@ class DefaultMcpSessionTests { transport.simulateIncomingMessage(notification); // Verify handler was called - assertThat(receivedParams.get()).isEqualTo(notificationParams); + assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams); } @Test diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java deleted file mode 100644 index 1362b32..0000000 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java +++ /dev/null @@ -1,154 +0,0 @@ -/* - * Copyright 2024-2024 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.ai.mcp.spec; - -import java.util.concurrent.atomic.AtomicReference; - -import com.fasterxml.jackson.databind.ObjectMapper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import reactor.core.publisher.Mono; -import reactor.test.StepVerifier; - -import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; -import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -class DefaultMcpTransportTests { - - private AbstractMcpTransport transport; - - private ObjectMapper objectMapper; - - @BeforeEach - void setUp() { - objectMapper = new ObjectMapper(); - transport = new AbstractMcpTransport(objectMapper) { - @Override - public Mono closeGracefully() { - return Mono.empty(); - } - }; - } - - @Test - void constructorShouldValidateObjectMapper() { - assertThatThrownBy(() -> new AbstractMcpTransport(null) { - @Override - public Mono closeGracefully() { - return null; - } - }).isInstanceOf(IllegalArgumentException.class).hasMessage("ObjectMapper must not be null"); - } - - @Test - void defaultConstructorShouldCreateValidInstance() { - AbstractMcpTransport defaultTransport = new AbstractMcpTransport() { - @Override - public Mono closeGracefully() { - return null; - } - }; - assertThat(defaultTransport.getObjectMapper()).isNotNull(); - } - - @Test - void customObjectMapperConstructorShouldCreateValidInstance() { - assertThat(transport.getObjectMapper()).isSameAs(objectMapper); - } - - @Test - void sendMessageShouldEmitToOutboundSink() { - JSONRPCRequest message = new JSONRPCRequest("2.0", "test", 1, null); - - AtomicReference receivedMessage = new AtomicReference<>(); - transport.getOutboundSink().asFlux().subscribe(receivedMessage::set); - - Mono result = transport.sendMessage(message); - - StepVerifier.create(result).verifyComplete(); - - assertThat(receivedMessage.get()).isNotNull().satisfies(msg -> { - assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0"); - assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test"); - }); - } - - @Test - void customInboundMessageHandlerShouldReceiveMessages() { - AtomicReference receivedMessage = new AtomicReference<>(); - transport.setInboundMessageHandler(receivedMessage::set); - transport.start(); - - JSONRPCRequest message = new JSONRPCRequest("2.0", "test", 1, null); - - transport.getInboundSink().tryEmitNext(message); - - assertThat(receivedMessage.get()).isNotNull().satisfies(msg -> { - assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0"); - assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test"); - }); - } - - @Test - void customErrorHandlerShouldReceiveErrors() { - AtomicReference receivedError = new AtomicReference<>(); - transport.setInboundErrorHandler(receivedError::set); - transport.start(); - - String errorMessage = "Test error"; - transport.getErrorSink().tryEmitNext(errorMessage); - - assertThat(receivedError.get()).isNotNull().isEqualTo(errorMessage); - } - - @Test - void startShouldInitializeMessageAndErrorHandling() { - AtomicReference receivedMessage = new AtomicReference<>(); - AtomicReference receivedError = new AtomicReference<>(); - - transport.setInboundMessageHandler(receivedMessage::set); - transport.setInboundErrorHandler(receivedError::set); - transport.start(); - - JSONRPCRequest message = new JSONRPCRequest("2.0", "test", 1, null); - - String errorMessage = "Test error"; - - transport.getInboundSink().tryEmitNext(message); - transport.getErrorSink().tryEmitNext(errorMessage); - - assertThat(receivedMessage.get()).isNotNull(); - assertThat(receivedError.get()).isEqualTo(errorMessage); - } - - @Test - void sendMessageShouldReturnErrorWhenSinkFails() { - JSONRPCRequest message = new JSONRPCRequest("2.0", "test", 1, null); - // Close the sink to simulate failure - transport.getOutboundSink().tryEmitComplete(); - - Mono result = transport.sendMessage(message); - - StepVerifier.create(result) - .verifyErrorSatisfies(error -> assertThat(error).isInstanceOf(RuntimeException.class) - .hasMessage("Failed to enqueue message")); - } - -}