diff --git a/src/main/java/spring/ai/mcp/client/McpAsyncSession.java b/src/main/java/spring/ai/mcp/client/McpAsyncSession.java index 875516c..7adb4f1 100644 --- a/src/main/java/spring/ai/mcp/client/McpAsyncSession.java +++ b/src/main/java/spring/ai/mcp/client/McpAsyncSession.java @@ -11,12 +11,14 @@ import reactor.core.publisher.Mono; import reactor.core.publisher.MonoSink; import spring.ai.mcp.spec.McpAsyncTransport; import spring.ai.mcp.spec.McpTransport; + +import org.springframework.util.Assert; + import spring.ai.mcp.spec.McpSchema; public class McpAsyncSession { - private final ConcurrentHashMap> pendingResponses - = new ConcurrentHashMap<>(); + private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); private final Duration requestTimeout; @@ -24,9 +26,12 @@ public class McpAsyncSession { private final McpAsyncTransport transport; - public McpAsyncSession(Duration requestTimeout, - ObjectMapper objectMapper, - McpAsyncTransport transport) { + public McpAsyncSession(Duration requestTimeout, ObjectMapper objectMapper, McpAsyncTransport transport) { + + Assert.notNull(requestTimeout, "Request timeout must not be null"); + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + Assert.notNull(transport, "Transport must not be null"); + this.requestTimeout = requestTimeout; this.objectMapper = objectMapper; this.transport = transport; @@ -37,7 +42,8 @@ public class McpAsyncSession { var sink = pendingResponses.remove(response.id()); if (sink == null) { System.out.println("Unexpected response for unkown id " + response.id()); - } else { + } + else { sink.success(response); } } @@ -60,34 +66,36 @@ public class McpAsyncSession { return Mono.create(sink -> { this.pendingResponses.put(requestId, sink); - McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, requestId, requestParams); + 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) { + // 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)); - } - } - }); + }).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 { @@ -95,6 +103,7 @@ public class McpAsyncSession { public McpError(Object error) { super(error.toString()); } + } public Mono sendNotification(String method) { @@ -102,8 +111,8 @@ public class McpAsyncSession { } public Mono sendNotification(String method, Map params) { - McpSchema.JSONRPCNotification - jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, method, params); + McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, + method, params); try { // TODO: make it non-blocking this.transport.sendMessage(jsonrpcNotification); @@ -118,4 +127,5 @@ public class McpAsyncSession { // TODO handle the timeout in transport return Mono.fromRunnable(this.transport::close); } + } diff --git a/src/main/java/spring/ai/mcp/client/McpSyncClient.java b/src/main/java/spring/ai/mcp/client/McpSyncClient.java index d02e489..370ac9f 100644 --- a/src/main/java/spring/ai/mcp/client/McpSyncClient.java +++ b/src/main/java/spring/ai/mcp/client/McpSyncClient.java @@ -4,6 +4,8 @@ import java.time.Duration; import spring.ai.mcp.spec.McpSchema; +import org.springframework.util.Assert; + public class McpSyncClient implements AutoCloseable { // TODO: Consider providing a client config to set this properly @@ -14,6 +16,7 @@ public class McpSyncClient implements AutoCloseable { McpAsyncClient delegate; public McpSyncClient(McpAsyncClient delegate) { + Assert.notNull(delegate, "McpAsyncClient must not be null"); this.delegate = delegate; } diff --git a/src/main/java/spring/ai/mcp/spec/McpSchema.java b/src/main/java/spring/ai/mcp/spec/McpSchema.java index b719772..2b0ee17 100644 --- a/src/main/java/spring/ai/mcp/spec/McpSchema.java +++ b/src/main/java/spring/ai/mcp/spec/McpSchema.java @@ -18,6 +18,7 @@ package spring.ai.mcp.spec; import java.util.List; import java.util.Map; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonSubTypes; @@ -213,6 +214,7 @@ public class McpSchema { * @param mimeType The MIME type of this resource, if known. */ @JsonInclude(JsonInclude.Include.NON_ABSENT) + @JsonIgnoreProperties(ignoreUnknown = true) public record Resource( // @formatter:off @JsonProperty("uri") String uri, @JsonProperty("name") String name, diff --git a/src/test/java/spring/ai/mcp/McpSyncClientTests.java b/src/test/java/spring/ai/mcp/McpSyncClientTests.java index 45feb4d..9000dfa 100644 --- a/src/test/java/spring/ai/mcp/McpSyncClientTests.java +++ b/src/test/java/spring/ai/mcp/McpSyncClientTests.java @@ -127,7 +127,7 @@ class McpSyncClientTests { .isNotNull() .satisfies(result -> { assertThat(result.content()).isNotNull(); - assertThat(result.isError()).isFalse(); + assertThat(result.isError()).isNull(); }); }