Factor out the McpSession into an inteface with DefaultMcpSession impl.

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

mcp/docs/mcp-class-diagram.puml

@@ -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<JSONRPCMessage>): void
-    +setErrorHandler(Consumer<String>): void
-    +sendMessage(JSONRPCMessage): Mono<Void>
-}
-
-interface McpAsyncSession {
-    +initialize(): Mono<InitializeResult>
-    +close(): void
-    +listTools(): Mono<ListToolsResult>
-    +callTool(String, Map<String, Object>): Mono<CallToolResult>
-    +createMessage(List<SamplingMessage>, ModelPreferences): Mono<CreateMessageResult>
-}
-
-interface McpClient {
-    +initialize(): InitializeResult
-    +listTools(): ListToolsResult
-    +callTool(String, Map<String, Object>): CallToolResult
-    +createMessage(List<SamplingMessage>, ModelPreferences): CreateMessageResult
-}
-
-' Classes
-class DefaultMcpTransport {
-    -objectMapper: ObjectMapper
-    -errorSink: Sinks.Many<String>
-    -inboundSink: Sinks.Many<JSONRPCMessage>
-    -outboundSink: Sinks.Many<JSONRPCMessage>
-    +DefaultMcpTransport()
-    +DefaultMcpTransport(ObjectMapper)
-}
-
-class McpAsyncClient {
-    -transport: McpTransport
-    -session: McpAsyncSession
-    +McpAsyncClient(McpTransport)
-    +initialize(): Mono<InitializeResult>
-    +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<String, String>
-}
-
-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

mcp/docs/mcp-sequence-diagram.puml

@@ -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

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());

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<Object, MonoSink<McpSchema.JSONRPCResponse>> 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 <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
+		String requestId = UUID.randomUUID().toString();
+
+		return Mono.<McpSchema.JSONRPCResponse>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<Void> sendNotification(String method, Map<String, Object> 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<Void> closeGracefully(Duration timeout) {
+		// TODO handle the timeout in transport
+		return Mono.fromRunnable(this.transport::close);
+	}
+
+}

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());
+	}
+
+}

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.
+ * 
+ * <p>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.</p>
  *
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
-public class McpSession {
+public interface McpSession {
 
-	private final ConcurrentHashMap<Object, MonoSink<McpSchema.JSONRPCResponse>> pendingResponses = new ConcurrentHashMap<>();
+    /**
+     * Sends a request to the model server and expects a response of type T.
+     * 
+     * <p>This method handles the request-response pattern where a response is expected
+     * from the server. The response type is determined by the provided TypeReference.</p>
+     *
+     * @param <T> 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
+     */
+    <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef);
 
-	private final Duration requestTimeout;
+    /**
-
+     * Sends a notification to the model server without parameters.
-	private final ObjectMapper objectMapper;
+     * 
-
+     * <p>This method implements the notification pattern where no response is expected
-	private final McpTransport transport;
+     * from the server. It's useful for fire-and-forget scenarios.</p>
-
+     *
-	public McpSession(Duration requestTimeout, ObjectMapper objectMapper, McpTransport transport) {
+     * @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
-		Assert.notNull(objectMapper, "The ObjectMapper can not be null");
+     */
-		Assert.notNull(requestTimeout, "The requstTimeout can not be null");
+	default Mono<Void> sendNotification(String method) {
-		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 <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
-		String requestId = UUID.randomUUID().toString();
-
-		return Mono.<McpSchema.JSONRPCResponse>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<Void> sendNotification(String method) {
 		return sendNotification(method, null);
 	}
 
-	public Mono<Void> sendNotification(String method, Map<String, Object> params) {
+    /**
-		McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
+     * Sends a notification to the model server with parameters.
-				method, params);
+     * 
-		try {
+     * <p>Similar to {@link #sendNotification(String)} but allows sending additional
-			this.transport.sendMessage(jsonrpcNotification);
+     * parameters with the notification.</p>
-		} catch (Exception e) {
+     *
-			return Mono.error(new McpError(e));
+     * @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 Mono.empty();
+     * @return a Mono that completes when the notification has been sent
-	}
+     */
-
+    Mono<Void> sendNotification(String method, Map<String, Object> params);
-	public Mono<Void> closeGracefully(Duration timeout) {
-		// TODO handle the timeout in transport
-		return Mono.fromRunnable(this.transport::close);
-	}
 
+    /**
+     * Gracefully closes the session with a specified timeout.
+     * 
+     * <p>This method attempts to close the session cleanly, waiting for any pending
+     * operations to complete up to the specified timeout duration.</p>
+     *
+     * @param timeout the maximum duration to wait for the session to close
+     * @return a Mono that completes when the session has been closed
+     */
+    Mono<Void> closeGracefully(Duration timeout);
 }

mcp/src/main/resources/application.properties

@@ -1,4 +0,0 @@
-spring.application.name=mcp
-spring.main.web-application-type=none
-
-spring.ai.openai.api-key=${OPENAI_API_KEY}