diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java index 596eda9..208eedf 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java @@ -44,7 +44,7 @@ public class McpSyncClient implements AutoCloseable { @Override public void close() { - this.delegate.closeGracefully(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS)).block(); + this.delegate.closeGracefully().block(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS)); } /** diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java index 9ec3a8f..a33c77c 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; import com.fasterxml.jackson.core.type.TypeReference; @@ -87,6 +88,11 @@ public class StdioServerTransport extends AbstractMcpTransport { @Override public void start() { + // Let's kick off the abstraction layer that will consume the logical messages + // pushed via sinks. The code that follows is actually feeding the sinks with + // data. + super.start(); + // Prepare command and environment List fullCommand = new ArrayList<>(); fullCommand.add(params.getCommand()); @@ -234,39 +240,43 @@ public class StdioServerTransport extends AbstractMcpTransport { this.isRunning = false; - return Mono.whenDelayError(Mono.fromRunnable(() -> { + return Mono.fromFuture(() -> { + System.out.println("Sending TERM to process"); + if (this.process != null) { + this.process.destroy(); + return process.onExit(); + } + else { + return CompletableFuture.failedFuture(new RuntimeException("Process not started")); + } + }).doOnNext(process -> { + if (process.exitValue() != 0) { + System.out.println("Process terminated with code " + process.exitValue()); + } + }).then(Mono.whenDelayError(Mono.fromRunnable(() -> { try { this.processErrorReader.close(); } catch (IOException e) { throw new RuntimeException(e); } - }).subscribeOn(errorScheduler), Mono.fromRunnable(() -> { + }), Mono.fromRunnable(() -> { try { this.processReader.close(); } catch (IOException e) { throw new RuntimeException(e); } - }).subscribeOn(inboundScheduler), Mono.fromRunnable(() -> { + }), Mono.fromRunnable(() -> { try { this.processWriter.close(); } catch (IOException e) { throw new RuntimeException(e); } - }).subscribeOn(outboundScheduler)) + }))) .then(Mono.whenDelayError(inboundScheduler.disposeGracefully(), outboundScheduler.disposeGracefully(), errorScheduler.disposeGracefully())) - .then(Mono.fromRunnable(() -> { - if (this.process != null) { - var process = this.process.destroyForcibly(); - if (process.exitValue() != 0) { - throw new RuntimeException("Failed to destroy process"); - } - } - })) - .then() .subscribeOn(Schedulers.boundedElastic()); } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java index 6c09c85..5e26319 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java @@ -29,6 +29,9 @@ import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage; * @author Christian Tzolov * @author Dariusz Jędrzejczyk */ +// TODO: The API is not finalized. Implementing SSE will dictate a more concrete +// structure and relationship between the base and the actual implementation including +// proper encapsulation and state management. public abstract class AbstractMcpTransport implements McpTransport { protected final ObjectMapper objectMapper; @@ -53,17 +56,9 @@ public abstract class AbstractMcpTransport implements McpTransport { Assert.notNull(objectMapper, "ObjectMapper must not be null"); this.objectMapper = objectMapper; - // TODO: consider the effects of buffering here -> the inter-process pipes are - // independent and the notifications can flood the client/server. - // Potentially, the interest in reading could be communicated from one party - // to the other so the Blocking IO Threads can pause consuming the stream - // buffers when there is no expectation for reading. - this.errorSink = Sinks.many().unicast().onBackpressureBuffer(); this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); - - this.handleIncomingMessages(); } public ObjectMapper getObjectMapper() { @@ -77,7 +72,6 @@ public abstract class AbstractMcpTransport implements McpTransport { private void handleIncomingErrors() { this.errorSink.asFlux().subscribe(e -> { this.errorHandler.accept(e); - // System.err.println(e); TODO: log the error }); } @@ -93,7 +87,7 @@ public abstract class AbstractMcpTransport implements McpTransport { return errorSink; } - public void setInboudMessageHandler(Consumer inboundMessageHandler) { + public void setInboundMessageHandler(Consumer inboundMessageHandler) { this.inboundMessageHandler = inboundMessageHandler; } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java index f754289..e8d3f4f 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java @@ -54,7 +54,7 @@ public class DefaultMcpSession implements McpSession { this.objectMapper = objectMapper; this.transport = transport; - this.transport.setInboudMessageHandler(message -> { + this.transport.setInboundMessageHandler(message -> { if (message instanceof McpSchema.JSONRPCResponse response) { var sink = pendingResponses.remove(response.id()); @@ -80,33 +80,27 @@ public class DefaultMcpSession implements McpSession { @Override public Mono sendRequest(String method, Object requestParams, TypeReference typeRef) { + // TODO: UUID API is blocking. Consider non-blocking alternatives to generate + // the ID. 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); - } + 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); + }); }).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")) { + if (typeRef.getType().equals(Void.class)) { s.complete(); } else { @@ -120,19 +114,17 @@ public class DefaultMcpSession implements McpSession { 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(); + return this.transport.sendMessage(jsonrpcNotification); } @Override - public Mono closeGracefully(Duration timeout) { - // TODO handle the timeout in transport - return Mono.fromRunnable(this.transport::close); + public Mono closeGracefully() { + return transport.closeGracefully(); + } + + @Override + public void close() { + transport.close(); } } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java index e78fa2d..fe0a56b 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java @@ -164,17 +164,24 @@ public class McpSchema { } // @formatter:on // Existing Enums and Base Types (from previous implementation) - public enum Role { + public enum Role {// @formatter:off - USER, ASSISTANT + @JsonProperty("user") USER, + @JsonProperty("assistant") ASSISTANT + }// @formatter:on - } + public enum LoggingLevel {// @formatter:off - public enum LoggingLevel { + @JsonProperty("debug") DEBUG, + @JsonProperty("info") INFO, + @JsonProperty("notice") NOTICE, + @JsonProperty("warning") WARNING, + @JsonProperty("error") ERROR, + @JsonProperty("critical") CRITICAL, + @JsonProperty("alert") ALERT, + @JsonProperty("emergency") EMERGENCY - DEBUG, INFO, NOTICE, WARNING, ERROR, CRITICAL, ALERT, EMERGENCY - - } + } // @formatter:on // --------------------------- // Resource Interfaces @@ -448,20 +455,22 @@ public class McpSchema { String systemPrompt, ContextInclusionStrategy includeContext, Double temperature, int maxTokens, List stopSequences, Map metadata) implements Request { - public enum ContextInclusionStrategy { - - NONE, THIS_SERVER, ALL_SERVERS - - } + public enum ContextInclusionStrategy {// @formatter:off + @JsonProperty("none") NONE, + @JsonProperty("this_server") THIS_SERVER, + @JsonProperty("all_server") ALL_SERVERS + } // @formatter:on } @JsonInclude(JsonInclude.Include.NON_ABSENT) public record CreateMessageResult(Role role, Content content, String model, StopReason stopReason) { public enum StopReason { - END_TURN, STOP_SEQUENCE, MAX_TOKENS - - } + // @formatter:off + @JsonProperty("end_turn") END_TURN, + @JsonProperty("stop_sequence") STOP_SEQUENCE, + @JsonProperty("max_tokens") MAX_TOKENS + } // @formatter:on } // --------------------------- diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSession.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSession.java index 7c0e48e..b1b6279 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSession.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSession.java @@ -81,15 +81,14 @@ public interface McpSession { 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 + * Closes the session and releases any associated resources asynchronously. + * @return a {@link Mono} that completes when the session has been closed. */ - Mono closeGracefully(Duration timeout); + Mono closeGracefully(); + + /** + * Closes the session and releases any associated resources. + */ + void close(); } diff --git a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java index 3d444bc..5fc6014 100644 --- a/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java +++ b/spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java @@ -91,7 +91,7 @@ public interface McpTransport { * @param inboundMessageHandler a consumer that processes incoming * {@link JSONRPCMessage}s */ - void setInboudMessageHandler(Consumer inboundMessageHandler); + void setInboundMessageHandler(Consumer inboundMessageHandler); /** * Sets the handler for processing transport-level errors. diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientTests.java index 70fc8d5..0ebfd6c 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientTests.java @@ -24,6 +24,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import reactor.test.StepVerifier; import org.springframework.ai.mcp.client.stdio.ServerParameters; import org.springframework.ai.mcp.client.stdio.StdioServerTransport; @@ -41,8 +42,9 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; * Unit tests for MCP Client Session functionality. * * @author Christian Tzolov - * @since 1.0.0 + * @author Dariusz Jędrzejczyk */ +@Timeout(15) // Giving extra time beyond the client timeout class McpAsyncClientTests { private McpAsyncClient mcpAsyncClient; @@ -61,15 +63,16 @@ class McpAsyncClientTests { assertThatCode(() -> { mcpAsyncClient = McpClient.async(new StdioServerTransport(stdioParams), TIMEOUT, new ObjectMapper()); - mcpAsyncClient.initialize(); + mcpAsyncClient.initialize().block(Duration.ofSeconds(10)); }).doesNotThrowAnyException(); } @AfterEach void tearDown() { - // if (mcpSyncClient != null) { - // assertThatCode(() -> mcpSyncClient.close()).doesNotThrowAnyException(); - // } + if (mcpAsyncClient != null) { + assertThatCode(() -> mcpAsyncClient.closeGracefully().block(Duration.ofSeconds(10))) + .doesNotThrowAnyException(); + } } @Test @@ -85,38 +88,34 @@ class McpAsyncClientTests { } @Test - @Timeout(15) // Giving extra time beyond the client timeout void testListTools() { - mcpAsyncClient.listTools(null).subscribe(result -> { + StepVerifier.create(mcpAsyncClient.listTools(null)).consumeNextWith(result -> { assertThat(result.tools()).isNotNull().isNotEmpty(); Tool firstTool = result.tools().get(0); assertThat(firstTool.name()).isNotNull(); assertThat(firstTool.description()).isNotNull(); - }); + }).verifyComplete(); } @Test - @Timeout(15) void testPing() { assertThatCode(() -> mcpAsyncClient.ping().block()).doesNotThrowAnyException(); } @Test - @Timeout(15) void testCallTool() { CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE)); - mcpAsyncClient.callTool(callToolRequest).subscribe(callToolResult -> { + StepVerifier.create(mcpAsyncClient.callTool(callToolRequest)).consumeNextWith(callToolResult -> { assertThat(callToolResult).isNotNull().satisfies(result -> { assertThat(result.content()).isNotNull(); assertThat(result.isError()).isNull(); }); - }); + }).verifyComplete(); } @Test - @Timeout(15) void testCallToolWithInvalidTool() { CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE)); @@ -124,15 +123,13 @@ class McpAsyncClientTests { } @Test - @Timeout(15) void testRootsListChanged() { assertThatCode(() -> mcpAsyncClient.sendRootsListChanged().block()).doesNotThrowAnyException(); } @Test - @Timeout(15) void testListResources() { - mcpAsyncClient.listResources(null).subscribe(resources -> { + StepVerifier.create(mcpAsyncClient.listResources(null)).consumeNextWith(resources -> { assertThat(resources).isNotNull().satisfies(result -> { assertThat(result.resources()).isNotNull(); @@ -142,7 +139,7 @@ class McpAsyncClientTests { assertThat(firstResource.name()).isNotNull(); } }); - }); + }).verifyComplete(); } @Test @@ -151,9 +148,8 @@ class McpAsyncClientTests { } @Test - @Timeout(15) void testListPrompts() { - mcpAsyncClient.listPrompts(null).subscribe(prompts -> { + StepVerifier.create(mcpAsyncClient.listPrompts(null)).consumeNextWith(prompts -> { assertThat(prompts).isNotNull().satisfies(result -> { assertThat(result.prompts()).isNotNull(); @@ -163,18 +159,19 @@ class McpAsyncClientTests { assertThat(firstPrompt.description()).isNotNull(); } }); - }); + }).verifyComplete(); } @Test - @Timeout(15) void testGetPrompt() { - mcpAsyncClient.getPrompt(new GetPromptRequest("simple_prompt", Map.of())).subscribe(prompt -> { - assertThat(prompt).isNotNull().satisfies(result -> { - assertThat(result.description()).isNotNull(); - assertThat(result.messages()).isNotEmpty(); - }); - }); + StepVerifier.create(mcpAsyncClient.getPrompt(new GetPromptRequest("simple_prompt", Map.of()))) + .consumeNextWith(prompt -> { + assertThat(prompt).isNotNull().satisfies(result -> { + assertThat(result.messages()).isNotEmpty(); + assertThat(result.messages()).hasSize(1); + }); + }) + .verifyComplete(); } } diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpSyncClientTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpSyncClientTests.java index 0a352cb..df81fab 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpSyncClientTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpSyncClientTests.java @@ -43,8 +43,9 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; * Unit tests for MCP Client Session functionality. * * @author Christian Tzolov - * @since 1.0.0 + * @author Dariusz Jędrzejczyk */ +@Timeout(15) // Giving extra time beyond the client timeout class McpSyncClientTests { private McpSyncClient mcpSyncClient; @@ -87,7 +88,6 @@ class McpSyncClientTests { } @Test - @Timeout(15) // Giving extra time beyond the client timeout void testListTools() { ListToolsResult tools = mcpSyncClient.listTools(null); @@ -101,7 +101,6 @@ class McpSyncClientTests { } @Test - @Timeout(15) // Giving extra time beyond the client timeout void testCallTools() { CallToolResult toolResult = mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4))); @@ -120,13 +119,11 @@ class McpSyncClientTests { } @Test - @Timeout(15) void testPing() { assertThatCode(() -> mcpSyncClient.ping()).doesNotThrowAnyException(); } @Test - @Timeout(15) void testCallTool() { CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE)); @@ -139,7 +136,6 @@ class McpSyncClientTests { } @Test - @Timeout(15) void testCallToolWithInvalidTool() { CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE)); @@ -147,13 +143,11 @@ class McpSyncClientTests { } @Test - @Timeout(15) void testRootsListChanged() { assertThatCode(() -> mcpSyncClient.sendRootsListChanged()).doesNotThrowAnyException(); } @Test - @Timeout(15) void testListResources() { ListResourcesResult resources = mcpSyncClient.listResources(null); diff --git a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java index 449ee3c..1362b32 100644 --- a/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java +++ b/spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java @@ -42,7 +42,7 @@ class DefaultMcpTransportTests { transport = new AbstractMcpTransport(objectMapper) { @Override public Mono closeGracefully() { - return null; + return Mono.empty(); } }; } @@ -93,7 +93,7 @@ class DefaultMcpTransportTests { @Test void customInboundMessageHandlerShouldReceiveMessages() { AtomicReference receivedMessage = new AtomicReference<>(); - transport.setInboudMessageHandler(receivedMessage::set); + transport.setInboundMessageHandler(receivedMessage::set); transport.start(); JSONRPCRequest message = new JSONRPCRequest("2.0", "test", 1, null); @@ -123,7 +123,7 @@ class DefaultMcpTransportTests { AtomicReference receivedMessage = new AtomicReference<>(); AtomicReference receivedError = new AtomicReference<>(); - transport.setInboudMessageHandler(receivedMessage::set); + transport.setInboundMessageHandler(receivedMessage::set); transport.setInboundErrorHandler(receivedError::set); transport.start();