modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Factor out the McpSession into an inteface with DefaultMcpSession impl.

Browse source
This commit is contained in:
Christian Tzolov 2024-12-10 14:03:11 +01:00
commit 7c9cf46e8a
4 changed files with 215 additions and 339 deletions
Download patch file Download diff file Expand all files Collapse all files

5
mcp/src/main/java/spring/ai/experimental/mcp/client/McpAsyncClient.java
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());

117
mcp/src/main/java/spring/ai/experimental/mcp/spec/DefaultMcpSession.java Normal file
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);
}
}

24
mcp/src/main/java/spring/ai/experimental/mcp/spec/McpError.java Normal file
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());
}
}

173
mcp/src/main/java/spring/ai/experimental/mcp/spec/McpSession.java
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);
}