Factor out the McpSession into an inteface with DefaultMcpSession impl.

This commit is contained in:
Christian Tzolov 2024-12-10 14:03:11 +01:00
commit 7c9cf46e8a
8 changed files with 215 additions and 339 deletions

View file

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

View file

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

View file

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

View file

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

View file

@ -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);
}
}

View file

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

View file

@ -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;
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 <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) {
/**
* Sends a notification to the model server without parameters.
*
* <p>This method implements the notification pattern where no response is expected
* from the server. It's useful for fire-and-forget scenarios.</p>
*
* @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<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,
method, params);
try {
this.transport.sendMessage(jsonrpcNotification);
} catch (Exception e) {
return Mono.error(new McpError(e));
}
return Mono.empty();
}
public Mono<Void> closeGracefully(Duration timeout) {
// TODO handle the timeout in transport
return Mono.fromRunnable(this.transport::close);
}
/**
* Sends a notification to the model server with parameters.
*
* <p>Similar to {@link #sendNotification(String)} but allows sending additional
* parameters with the notification.</p>
*
* @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<Void> sendNotification(String method, Map<String, Object> params);
/**
* 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);
}

View file

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