From 7c9cf46e8abd242b5b88c58a24d1193c07e7bd01 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 10 Dec 2024 14:03:11 +0100 Subject: [PATCH] Factor out the McpSession into an inteface with DefaultMcpSession impl. --- mcp/README.md | 2 +- mcp/docs/mcp-class-diagram.puml | 118 ------------ mcp/docs/mcp-sequence-diagram.puml | 105 ----------- .../mcp/client/McpAsyncClient.java | 5 +- .../mcp/spec/DefaultMcpSession.java | 117 ++++++++++++ .../ai/experimental/mcp/spec/McpError.java | 24 +++ .../ai/experimental/mcp/spec/McpSession.java | 173 +++++++----------- mcp/src/main/resources/application.properties | 4 - 8 files changed, 212 insertions(+), 336 deletions(-) delete mode 100644 mcp/docs/mcp-class-diagram.puml delete mode 100644 mcp/docs/mcp-sequence-diagram.puml create mode 100644 mcp/src/main/java/spring/ai/experimental/mcp/spec/DefaultMcpSession.java create mode 100644 mcp/src/main/java/spring/ai/experimental/mcp/spec/McpError.java delete mode 100644 mcp/src/main/resources/application.properties diff --git a/mcp/README.md b/mcp/README.md index 88e0cc3..1263060 100644 --- a/mcp/README.md +++ b/mcp/README.md @@ -102,7 +102,7 @@ The SDK follows a layered architecture: - **McpAsyncClient**: Async implementation using Project Reactor - **McpSyncClient**: Synchronous wrapper around the async client - **McpTransport**: Transport layer interface -- **DefaultMcpTransport**: Base transport implementation +- **AbstracttMcpTransport**: Abstract transport implementation - **StdioServerTransport**: Stdio-based server communication ### Key Interactions diff --git a/mcp/docs/mcp-class-diagram.puml b/mcp/docs/mcp-class-diagram.puml deleted file mode 100644 index 5c0d939..0000000 --- a/mcp/docs/mcp-class-diagram.puml +++ /dev/null @@ -1,118 +0,0 @@ -@startuml MCP Class Diagram - -' Style configurations -skinparam class { - BackgroundColor White - ArrowColor Black - BorderColor Black -} - -' Interfaces -interface McpTransport { - +start(): void - +close(): void - +setInboudMessageHandler(Consumer): void - +setErrorHandler(Consumer): void - +sendMessage(JSONRPCMessage): Mono -} - -interface McpAsyncSession { - +initialize(): Mono - +close(): void - +listTools(): Mono - +callTool(String, Map): Mono - +createMessage(List, ModelPreferences): Mono -} - -interface McpClient { - +initialize(): InitializeResult - +listTools(): ListToolsResult - +callTool(String, Map): CallToolResult - +createMessage(List, ModelPreferences): CreateMessageResult -} - -' Classes -class DefaultMcpTransport { - -objectMapper: ObjectMapper - -errorSink: Sinks.Many - -inboundSink: Sinks.Many - -outboundSink: Sinks.Many - +DefaultMcpTransport() - +DefaultMcpTransport(ObjectMapper) -} - -class McpAsyncClient { - -transport: McpTransport - -session: McpAsyncSession - +McpAsyncClient(McpTransport) - +initialize(): Mono - +close(): void -} - -class McpSyncClient { - -asyncClient: McpAsyncClient - +McpSyncClient(McpTransport) - +initialize(): InitializeResult - +close(): void -} - -class StdioServerTransport { - -process: Process - -writer: BufferedWriter - -reader: BufferedReader - +StdioServerTransport(ServerParameters) -} - -class ServerParameters { - -command: String[] - -workingDirectory: File - -environment: Map -} - -class ServerParametersParser { - +{static} parse(String): ServerParameters -} - -class Assert { - +{static} notNull(Object, String): void - +{static} hasText(String, String): void -} - -' Relationships -McpTransport <|.. DefaultMcpTransport -DefaultMcpTransport <|-- StdioServerTransport - -McpClient <|.. McpSyncClient -McpAsyncSession <|.. McpAsyncClient - -McpSyncClient o-- McpAsyncClient -McpAsyncClient o-- McpTransport -McpAsyncClient o-- McpAsyncSession - -StdioServerTransport o-- ServerParameters -ServerParameters <.. ServerParametersParser : creates - -' Package organization -package "spring.ai.mcp.spec" { - ' McpTransport - ' McpAsyncSession - ' DefaultMcpTransport -} - -package "spring.ai.mcp.client" { - ' McpClient - ' McpAsyncClient - ' McpSyncClient - - package "stdio" { - ' StdioServerTransport - ' ServerParameters - ' ServerParametersParser - } - - package "util" { - ' Assert - } -} - -@enduml diff --git a/mcp/docs/mcp-sequence-diagram.puml b/mcp/docs/mcp-sequence-diagram.puml deleted file mode 100644 index cbd09af..0000000 --- a/mcp/docs/mcp-sequence-diagram.puml +++ /dev/null @@ -1,105 +0,0 @@ -@startuml MCP Sequence Diagram - -participant McpClient -participant McpSyncClient -participant McpAsyncClient -participant DefaultMcpTransport -participant StdioServerTransport - -== Initialization == - -McpClient -> McpSyncClient: new McpSyncClient(transport) -activate McpSyncClient - -McpSyncClient -> McpAsyncClient: new McpAsyncClient(transport) -activate McpAsyncClient - -McpAsyncClient -> DefaultMcpTransport: start() -activate DefaultMcpTransport -DefaultMcpTransport --> McpAsyncClient: started -deactivate DefaultMcpTransport - -McpAsyncClient --> McpSyncClient: created -deactivate McpAsyncClient - -McpSyncClient --> McpClient: created -deactivate McpSyncClient - -== Message Flow == - -McpClient -> McpSyncClient: initialize() -activate McpSyncClient - -McpSyncClient -> McpAsyncClient: initialize() -activate McpAsyncClient - -McpAsyncClient -> DefaultMcpTransport: sendMessage(initRequest) -activate DefaultMcpTransport - -DefaultMcpTransport -> StdioServerTransport: sendMessage(initRequest) -activate StdioServerTransport - -StdioServerTransport --> DefaultMcpTransport: response -deactivate StdioServerTransport - -DefaultMcpTransport --> McpAsyncClient: InitializeResult -deactivate DefaultMcpTransport - -McpAsyncClient --> McpSyncClient: InitializeResult -deactivate McpAsyncClient - -McpSyncClient --> McpClient: InitializeResult -deactivate McpSyncClient - -== Tool Execution == - -McpClient -> McpSyncClient: callTool(name, args) -activate McpSyncClient - -McpSyncClient -> McpAsyncClient: callTool(name, args) -activate McpAsyncClient - -McpAsyncClient -> DefaultMcpTransport: sendMessage(toolRequest) -activate DefaultMcpTransport - -DefaultMcpTransport -> StdioServerTransport: sendMessage(toolRequest) -activate StdioServerTransport - -StdioServerTransport --> DefaultMcpTransport: response -deactivate StdioServerTransport - -DefaultMcpTransport --> McpAsyncClient: CallToolResult -deactivate DefaultMcpTransport - -McpAsyncClient --> McpSyncClient: CallToolResult -deactivate McpAsyncClient - -McpSyncClient --> McpClient: CallToolResult -deactivate McpSyncClient - -== Cleanup == - -McpClient -> McpSyncClient: close() -activate McpSyncClient - -McpSyncClient -> McpAsyncClient: close() -activate McpAsyncClient - -McpAsyncClient -> DefaultMcpTransport: close() -activate DefaultMcpTransport - -DefaultMcpTransport -> StdioServerTransport: close() -activate StdioServerTransport -StdioServerTransport --> DefaultMcpTransport: closed -deactivate StdioServerTransport - -DefaultMcpTransport --> McpAsyncClient: closed -deactivate DefaultMcpTransport - -McpAsyncClient --> McpSyncClient: closed -deactivate McpAsyncClient - -McpSyncClient --> McpClient: closed -deactivate McpSyncClient - -@enduml diff --git a/mcp/src/main/java/spring/ai/experimental/mcp/client/McpAsyncClient.java b/mcp/src/main/java/spring/ai/experimental/mcp/client/McpAsyncClient.java index eae2b56..9dce57e 100644 --- a/mcp/src/main/java/spring/ai/experimental/mcp/client/McpAsyncClient.java +++ b/mcp/src/main/java/spring/ai/experimental/mcp/client/McpAsyncClient.java @@ -5,15 +5,16 @@ import java.time.Duration; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import reactor.core.publisher.Mono; -import spring.ai.experimental.mcp.spec.McpSession; +import spring.ai.experimental.mcp.spec.DefaultMcpSession; import spring.ai.experimental.mcp.spec.McpTransport; +import spring.ai.experimental.mcp.spec.McpError; import spring.ai.experimental.mcp.spec.McpSchema; /** * @author Dariusz Jędrzejczyk * @author Christian Tzolov */ -public class McpAsyncClient extends McpSession { +public class McpAsyncClient extends DefaultMcpSession { public McpAsyncClient(McpTransport transport) { this(transport, Duration.ofSeconds(10), new ObjectMapper()); diff --git a/mcp/src/main/java/spring/ai/experimental/mcp/spec/DefaultMcpSession.java b/mcp/src/main/java/spring/ai/experimental/mcp/spec/DefaultMcpSession.java new file mode 100644 index 0000000..d736d9b --- /dev/null +++ b/mcp/src/main/java/spring/ai/experimental/mcp/spec/DefaultMcpSession.java @@ -0,0 +1,117 @@ +package spring.ai.experimental.mcp.spec; + +import java.time.Duration; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import reactor.core.publisher.Mono; +import reactor.core.publisher.MonoSink; +import spring.ai.experimental.mcp.client.util.Assert; + +/** + * Implementation of the MCP client session. + * + * @author Christian Tzolov + * @author Dariusz Jędrzejczyk + */ +public class DefaultMcpSession implements McpSession { + + private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); + + private final Duration requestTimeout; + + private final ObjectMapper objectMapper; + + private final McpTransport transport; + + public DefaultMcpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport) { + + 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"); + + this.requestTimeout = requestTimeout; + this.objectMapper = objectMapper; + this.transport = transport; + + this.transport.setInboudMessageHandler(message -> { + switch (message) { + case McpSchema.JSONRPCResponse response -> { + var sink = pendingResponses.remove(response.id()); + if (sink == null) { + System.out.println("Unexpected response for unkown id " + response.id()); + } else { + sink.success(response); + } + } + case McpSchema.JSONRPCRequest request -> { + System.out.println("Client does not yet support server requests"); + } + case McpSchema.JSONRPCNotification notification -> { + System.out.println("Notifications not yet supported"); + } + } + }); + + this.transport.setInboundErrorHandler(error -> System.out.println("Error received " + error)); + + this.transport.start(); + } + + @Override + public Mono sendRequest(String method, Object requestParams, TypeReference typeRef) { + String requestId = UUID.randomUUID().toString(); + + return Mono.create(sink -> { + this.pendingResponses.put(requestId, sink); + McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, + requestId, requestParams); + try { + // TODO: This is non-blocking, but it's actually a synchronous call, + // perhaps there's no need to make it return Mono? + this.transport.sendMessage(jsonrpcRequest) + // TODO: It's most efficient to create a dedicated + // Subscriber here + .subscribe(v -> { + }, e -> { + this.pendingResponses.remove(requestId); + sink.error(e); + }); + } catch (Exception e) { + sink.error(e); + } + }).timeout(this.requestTimeout).handle((jsonRpcResponse, s) -> { + if (jsonRpcResponse.error() != null) { + s.error(new McpError(jsonRpcResponse.error())); + } else { + if (typeRef.getType().getTypeName().equals("java.lang.Void")) { + s.complete(); + } else { + s.next(this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef)); + } + } + }); + } + + @Override + public Mono sendNotification(String method, Map params) { + McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, + method, params); + try { + this.transport.sendMessage(jsonrpcNotification); + } catch (Exception e) { + return Mono.error(new McpError(e)); + } + return Mono.empty(); + } + + @Override + public Mono closeGracefully(Duration timeout) { + // TODO handle the timeout in transport + return Mono.fromRunnable(this.transport::close); + } + +} diff --git a/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpError.java b/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpError.java new file mode 100644 index 0000000..408b55b --- /dev/null +++ b/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpError.java @@ -0,0 +1,24 @@ +/* +* 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 spring.ai.experimental.mcp.spec; + +public class McpError extends RuntimeException { + + public McpError(Object error) { + super(error.toString()); + } + +} \ No newline at end of file diff --git a/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpSession.java b/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpSession.java index 894a0ee..f5336a0 100644 --- a/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpSession.java +++ b/mcp/src/main/java/spring/ai/experimental/mcp/spec/McpSession.java @@ -1,126 +1,87 @@ +/* + * 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 spring.ai.experimental.mcp.spec; import java.time.Duration; import java.util.Map; -import java.util.UUID; -import java.util.concurrent.ConcurrentHashMap; import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; import reactor.core.publisher.Mono; -import reactor.core.publisher.MonoSink; -import spring.ai.experimental.mcp.client.util.Assert; /** - * Implementation of the MCP client session. + * Represents a Model Control Protocol (MCP) session that handles communication between + * clients and the server. This interface provides methods for sending requests + * and notifications, as well as managing the session lifecycle. * + *

The session operates asynchronously using Project Reactor's {@link Mono} type + * for non-blocking operations. It supports both request-response patterns and + * one-way notifications.

+ * * @author Christian Tzolov * @author Dariusz Jędrzejczyk */ -public class McpSession { +public interface McpSession { - private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); + /** + * Sends a request to the model server and expects a response of type T. + * + *

This method handles the request-response pattern where a response is expected + * from the server. The response type is determined by the provided TypeReference.

+ * + * @param the type of the expected response + * @param method the name of the method to be called on the server + * @param requestParams the parameters to be sent with the request + * @param typeRef the TypeReference describing the expected response type + * @return a Mono that will emit the response when received + */ + Mono sendRequest(String method, Object requestParams, TypeReference typeRef); - private final Duration requestTimeout; - - private final ObjectMapper objectMapper; - - private final McpTransport transport; - - public McpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport) { - - 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"); - - this.requestTimeout = requestTimeout; - this.objectMapper = objectMapper; - this.transport = transport; - - this.transport.setInboudMessageHandler(message -> { - switch (message) { - case McpSchema.JSONRPCResponse response -> { - var sink = pendingResponses.remove(response.id()); - if (sink == null) { - System.out.println("Unexpected response for unkown id " + response.id()); - } else { - sink.success(response); - } - } - case McpSchema.JSONRPCRequest request -> { - System.out.println("Client does not yet support server requests"); - } - case McpSchema.JSONRPCNotification notification -> { - System.out.println("Notifications not yet supported"); - } - } - }); - - this.transport.setInboundErrorHandler(error -> System.out.println("Error received " + error)); - - this.transport.start(); - } - - public Mono sendRequest(String method, Object requestParams, TypeReference typeRef) { - String requestId = UUID.randomUUID().toString(); - - return Mono.create(sink -> { - this.pendingResponses.put(requestId, sink); - McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, - requestId, requestParams); - try { - // TODO: This is non-blocking, but it's actually a synchronous call, - // perhaps there's no need to make it return Mono? - this.transport.sendMessage(jsonrpcRequest) - // TODO: It's most efficient to create a dedicated - // Subscriber here - .subscribe(v -> { - }, e -> { - this.pendingResponses.remove(requestId); - sink.error(e); - }); - } catch (Exception e) { - sink.error(e); - } - }).timeout(this.requestTimeout).handle((jsonRpcResponse, s) -> { - if (jsonRpcResponse.error() != null) { - s.error(new McpError(jsonRpcResponse.error())); - } else { - if (typeRef.getType().getTypeName().equals("java.lang.Void")) { - s.complete(); - } else { - s.next(this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef)); - } - } - }); - } - - public static class McpError extends RuntimeException { - - public McpError(Object error) { - super(error.toString()); - } - - } - - public Mono sendNotification(String method) { + /** + * Sends a notification to the model server without parameters. + * + *

This method implements the notification pattern where no response is expected + * from the server. It's useful for fire-and-forget scenarios.

+ * + * @param method the name of the notification method to be called on the server + * @return a Mono that completes when the notification has been sent + */ + default Mono sendNotification(String method) { return sendNotification(method, null); } - public Mono sendNotification(String method, Map params) { - McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, - method, params); - try { - this.transport.sendMessage(jsonrpcNotification); - } catch (Exception e) { - return Mono.error(new McpError(e)); - } - return Mono.empty(); - } - - public Mono closeGracefully(Duration timeout) { - // TODO handle the timeout in transport - return Mono.fromRunnable(this.transport::close); - } + /** + * Sends a notification to the model server with parameters. + * + *

Similar to {@link #sendNotification(String)} but allows sending additional + * parameters with the notification.

+ * + * @param method the name of the notification method to be called on the server + * @param params a map of parameters to be sent with the notification + * @return a Mono that completes when the notification has been sent + */ + Mono sendNotification(String method, Map params); + /** + * Gracefully closes the session with a specified timeout. + * + *

This method attempts to close the session cleanly, waiting for any pending + * operations to complete up to the specified timeout duration.

+ * + * @param timeout the maximum duration to wait for the session to close + * @return a Mono that completes when the session has been closed + */ + Mono closeGracefully(Duration timeout); } diff --git a/mcp/src/main/resources/application.properties b/mcp/src/main/resources/application.properties deleted file mode 100644 index 863de4f..0000000 --- a/mcp/src/main/resources/application.properties +++ /dev/null @@ -1,4 +0,0 @@ -spring.application.name=mcp -spring.main.web-application-type=none - -spring.ai.openai.api-key=${OPENAI_API_KEY} \ No newline at end of file