diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java b/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java index 8846774..31c1780 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java @@ -77,7 +77,7 @@ public class StdioServerTransport implements McpTransport { * Creates a new StdioServerTransport with a default ObjectMapper and System streams. */ public StdioServerTransport() { - this(new ObjectMapper(), System.in, System.out); + this(new ObjectMapper()); } /** @@ -86,26 +86,15 @@ public class StdioServerTransport implements McpTransport { * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization */ public StdioServerTransport(ObjectMapper objectMapper) { - this(objectMapper, System.in, System.out); - } - /** - * Creates a new StdioServerTransport with the specified ObjectMapper and streams. - * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization - * @param inputStream The input stream to read from - * @param outputStream The output stream to write to - */ - public StdioServerTransport(ObjectMapper objectMapper, InputStream inputStream, OutputStream outputStream) { Assert.notNull(objectMapper, "The ObjectMapper can not be null"); - Assert.notNull(inputStream, "The InputStream can not be null"); - Assert.notNull(outputStream, "The OutputStream can not be null"); this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.objectMapper = objectMapper; - this.inputStream = inputStream; - this.outputStream = outputStream; + this.inputStream = System.in; + this.outputStream = System.out; // Use bounded schedulers for better resource management this.inboundScheduler = Schedulers.newBoundedElastic(1, 1, "inbound"); @@ -248,50 +237,31 @@ public class StdioServerTransport implements McpTransport { return Mono.fromRunnable(() -> { isClosing = true; logger.debug("Initiating graceful shutdown"); - }) - // .then(Mono.defer(() -> { - // // First complete the sinks to stop processing - // inboundSink.tryEmitComplete(); - // outboundSink.tryEmitComplete(); - // return Mono.delay(Duration.ofMillis(100)); - // })) - .then(Mono.fromRunnable(() -> { - try { - // Dispose schedulers first - inboundScheduler.dispose(); - outboundScheduler.dispose(); + }).then(Mono.defer(() -> { + // First complete the sinks to stop processing + inboundSink.tryEmitComplete(); + outboundSink.tryEmitComplete(); + return Mono.delay(Duration.ofMillis(100)); + })).then(Mono.fromRunnable(() -> { + try { + // Dispose schedulers first + inboundScheduler.dispose(); + outboundScheduler.dispose(); - // Wait for schedulers to terminate - if (!inboundScheduler.isDisposed()) { - inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5)); - } - if (!outboundScheduler.isDisposed()) { - outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5)); - } - - // Only after schedulers are disposed, close the streams - try { - if (inputStream != System.in) { - inputStream.close(); - } - if (outputStream != System.out) { - outputStream.flush(); - outputStream.close(); - } - } - catch (IOException e) { - // Log but don't throw since we're shutting down - logger.debug("Error closing streams during shutdown", e); - } - - logger.info("Graceful shutdown completed"); + // Wait for schedulers to terminate + if (!inboundScheduler.isDisposed()) { + inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5)); } - catch (Exception e) { - logger.error("Error during graceful shutdown", e); + if (!outboundScheduler.isDisposed()) { + outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5)); } - })) - .then() - .subscribeOn(Schedulers.boundedElastic()); + + logger.info("Graceful shutdown completed"); + } + catch (Exception e) { + logger.error("Error during graceful shutdown", e); + } + })).then().subscribeOn(Schedulers.boundedElastic()); } @Override diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpAsyncServerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpAsyncServerTests.java index 9a0185c..9d5c80b 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpAsyncServerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpAsyncServerTests.java @@ -16,8 +16,6 @@ package org.springframework.ai.mcp.server; -import com.fasterxml.jackson.databind.ObjectMapper; - import org.springframework.ai.mcp.server.transport.StdioServerTransport; import org.springframework.ai.mcp.spec.McpTransport; @@ -30,7 +28,7 @@ class StdioMcpAsyncServerTests extends AbstractMcpAsyncServerTests { @Override protected McpTransport createMcpTransport() { - return new StdioServerTransport(new ObjectMapper()); + return new StdioServerTransport(); } } diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpSyncServerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpSyncServerTests.java index 745264a..a90e4de 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpSyncServerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpSyncServerTests.java @@ -16,11 +16,6 @@ package org.springframework.ai.mcp.server; -import java.io.PipedInputStream; -import java.io.PipedOutputStream; - -import com.fasterxml.jackson.databind.ObjectMapper; - import org.springframework.ai.mcp.server.transport.StdioServerTransport; import org.springframework.ai.mcp.spec.McpTransport; @@ -31,53 +26,9 @@ import org.springframework.ai.mcp.spec.McpTransport; */ class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests { - private PipedInputStream testInput; - - private PipedOutputStream testOutput; - - private PipedOutputStream writeToInput; - - private PipedInputStream readFromOutput; - @Override protected McpTransport createMcpTransport() { - try { - testInput = new PipedInputStream(); - writeToInput = new PipedOutputStream(testInput); - testOutput = new PipedOutputStream(); - readFromOutput = new PipedInputStream(testOutput); - return new StdioServerTransport(new ObjectMapper(), testInput, testOutput); - } - catch (Exception e) { - throw new RuntimeException("Failed to initialize test streams", e); - } - } - - @Override - protected void onStart() { - // No special setup needed for stdio transport - } - - @Override - protected void onClose() { - try { - if (testInput != null) { - testInput.close(); - } - if (testOutput != null) { - testOutput.close(); - } - if (writeToInput != null) { - writeToInput.close(); - } - if (readFromOutput != null) { - readFromOutput.close(); - } - } - catch (Exception e) { - // Log but don't throw since this is cleanup - e.printStackTrace(); - } + return new StdioServerTransport(); } } diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/transport/StdioServerTransportTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/transport/StdioServerTransportTests.java index 34e549b..ba0446d 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/transport/StdioServerTransportTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/transport/StdioServerTransportTests.java @@ -92,7 +92,7 @@ class StdioServerTransportTests { testIn = new ByteArrayInputStream((jsonMessage + "\n").getBytes(StandardCharsets.UTF_8)); // Create transport with test streams - transport = new StdioServerTransport(objectMapper, testIn, testOutPrintStream); + transport = new StdioServerTransport(objectMapper); // Parse expected message McpSchema.JSONRPCRequest expected = objectMapper.readValue(jsonMessage, McpSchema.JSONRPCRequest.class); @@ -109,9 +109,9 @@ class StdioServerTransportTests { @Disabled void shouldHandleOutgoingMessages() throws Exception { // Create transport with test streams - transport = new StdioServerTransport(objectMapper, new BlockingInputStream(), testOutPrintStream); - // transport = new StdioServerTransport(objectMapper, new ByteArrayInputStream(new - // byte[0]), testOutPrintStream); + transport = new StdioServerTransport(objectMapper); + // transport = new StdioServerTransport(objectMapper, new BlockingInputStream(), + // testOutPrintStream); // Create test messages JSONRPCRequest initMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "init", "init-id", @@ -137,7 +137,7 @@ class StdioServerTransportTests { @Test void shouldWaitForProcessorsBeforeSendingMessage() { // Create transport with test streams - transport = new StdioServerTransport(objectMapper, new ByteArrayInputStream(new byte[0]), testOutPrintStream); + transport = new StdioServerTransport(objectMapper); // Create test message JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test", "test-id", @@ -154,7 +154,7 @@ class StdioServerTransportTests { @Test void shouldCloseGracefully() { // Create transport with test streams - transport = new StdioServerTransport(objectMapper, new ByteArrayInputStream(new byte[0]), testOutPrintStream); + transport = new StdioServerTransport(objectMapper); // Create test message JSONRPCRequest initMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "init", "init-id",