diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java index 7e1d259..e7c07f4 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java @@ -349,9 +349,9 @@ public class McpAsyncServer { private DefaultMcpSession.RequestHandler resourcesListRequestHandler() { return params -> { - McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); + // McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, + // new TypeReference() { + // }); var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList(); @@ -361,9 +361,9 @@ public class McpAsyncServer { private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() { return params -> { - McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); + // McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, + // new TypeReference() { + // }); return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null)); }; 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 31c1780..675737f 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 @@ -23,6 +23,7 @@ import java.io.InputStreamReader; import java.io.OutputStream; import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.util.concurrent.Executors; import java.util.function.Function; import com.fasterxml.jackson.core.type.TypeReference; @@ -97,8 +98,8 @@ public class StdioServerTransport implements McpTransport { this.outputStream = System.out; // Use bounded schedulers for better resource management - this.inboundScheduler = Schedulers.newBoundedElastic(1, 1, "inbound"); - this.outboundScheduler = Schedulers.newBoundedElastic(1, 1, "outbound"); + this.inboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "inbound"); + this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "outbound"); } @Override @@ -139,38 +140,43 @@ public class StdioServerTransport implements McpTransport { private void startInboundProcessing() { this.inboundScheduler.schedule(() -> { inboundReady.tryEmitValue(null); - try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) { - String line; - while (!isClosing && (line = reader.readLine()) != null) { + BufferedReader reader = null; + try { + reader = new BufferedReader(new InputStreamReader(inputStream)); + while (!isClosing) { try { - JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line); - if (!this.inboundSink.tryEmitNext(message).isSuccess()) { + String line = reader.readLine(); + if (line == null || isClosing) { + break; + } + + try { + JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line); + if (!this.inboundSink.tryEmitNext(message).isSuccess()) { + if (!isClosing) { + logger.error("Failed to enqueue message"); + } + break; + } + } + catch (Exception e) { if (!isClosing) { - logger.error("Failed to enqueue message"); + logger.error("Error processing inbound message", e); } break; } } - catch (Exception e) { + catch (IOException e) { if (!isClosing) { - logger.error("Error processing inbound message", e); + logger.error("Error reading from stdin", e); } break; } } } - catch (IOException e) { - // Check isClosing before the error occurs to properly categorize it - boolean wasClosing = isClosing; - isClosing = true; - if (!wasClosing && e.getMessage().equals("Pipe closed")) { - logger.debug("Stream closed during shutdown", e); - } - else if (!wasClosing) { - logger.error("Error reading from stdin", e); - } - else { - logger.debug("Stream error during shutdown", e); + catch (Exception e) { + if (!isClosing) { + logger.error("Error in inbound processing", e); } } finally { @@ -234,6 +240,7 @@ public class StdioServerTransport implements McpTransport { @Override public Mono closeGracefully() { + return Mono.fromRunnable(() -> { isClosing = true; logger.debug("Initiating graceful shutdown"); @@ -244,18 +251,10 @@ public class StdioServerTransport implements McpTransport { return Mono.delay(Duration.ofMillis(100)); })).then(Mono.fromRunnable(() -> { try { - // Dispose schedulers first + // Dispose schedulers with longer timeout 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)); - } - logger.info("Graceful shutdown completed"); } catch (Exception e) { diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java index 15ef810..8807d40 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java @@ -66,6 +66,7 @@ public abstract class AbstractMcpSyncServerTests { @BeforeEach void setUp() { + // onStart(); } @AfterEach diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpAsyncServerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpAsyncServerTests.java index b9a2b77..3d8336b 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpAsyncServerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpAsyncServerTests.java @@ -17,6 +17,8 @@ package org.springframework.ai.mcp.server; import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Timeout; + import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; import org.springframework.web.reactive.function.server.RouterFunctions; @@ -31,6 +33,7 @@ import org.springframework.ai.mcp.spec.McpTransport; * * @author Christian Tzolov */ +@Timeout(15) // Giving extra time beyond the client timeout class SseMcpAsyncServerTests extends AbstractMcpAsyncServerTests { private static final int PORT = 8181; diff --git a/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpSyncServerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpSyncServerTests.java index 2d92574..c742969 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpSyncServerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpSyncServerTests.java @@ -17,6 +17,8 @@ package org.springframework.ai.mcp.server; import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Timeout; + import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; import org.springframework.web.reactive.function.server.RouterFunctions; @@ -31,6 +33,7 @@ import org.springframework.ai.mcp.spec.McpTransport; * * @author Christian Tzolov */ +@Timeout(15) // Giving extra time beyond the client timeout class SseMcpSyncServerTests extends AbstractMcpSyncServerTests { private static final int PORT = 8182; 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 9d5c80b..8d8a31d 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,6 +16,8 @@ package org.springframework.ai.mcp.server; +import org.junit.jupiter.api.Timeout; + import org.springframework.ai.mcp.server.transport.StdioServerTransport; import org.springframework.ai.mcp.spec.McpTransport; @@ -24,6 +26,7 @@ import org.springframework.ai.mcp.spec.McpTransport; * * @author Christian Tzolov */ +@Timeout(15) // Giving extra time beyond the client timeout class StdioMcpAsyncServerTests extends AbstractMcpAsyncServerTests { @Override 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 a90e4de..2e02f0d 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,6 +16,8 @@ package org.springframework.ai.mcp.server; +import org.junit.jupiter.api.Timeout; + import org.springframework.ai.mcp.server.transport.StdioServerTransport; import org.springframework.ai.mcp.spec.McpTransport; @@ -24,6 +26,7 @@ import org.springframework.ai.mcp.spec.McpTransport; * * @author Christian Tzolov */ +@Timeout(15) // Giving extra time beyond the client timeout class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests { @Override 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 ba0446d..c6ce443 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 @@ -16,7 +16,6 @@ package org.springframework.ai.mcp.server.transport; -import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.InputStream; import java.io.PrintStream; @@ -51,8 +50,6 @@ class StdioServerTransportTests { private ByteArrayOutputStream testOut; - private ByteArrayInputStream testIn; - private ByteArrayOutputStream testErr; private PrintStream testOutPrintStream; @@ -89,7 +86,6 @@ class StdioServerTransportTests { void shouldHandleIncomingMessages() throws Exception { // Prepare test input String jsonMessage = "{\"jsonrpc\":\"2.0\",\"method\":\"test\",\"params\":{},\"id\":1}"; - testIn = new ByteArrayInputStream((jsonMessage + "\n").getBytes(StandardCharsets.UTF_8)); // Create transport with test streams transport = new StdioServerTransport(objectMapper); diff --git a/pom.xml b/pom.xml index 4a1823b..383a2d2 100644 --- a/pom.xml +++ b/pom.xml @@ -174,6 +174,7 @@ false +