diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java index 73522ae..60a8d7c 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java @@ -258,16 +258,13 @@ public class StdioClientTransport implements McpClientTransport { @Override public Mono sendMessage(JSONRPCMessage message) { - if (this.outboundSink.tryEmitNext(message).isSuccess()) { - // TODO: essentially we could reschedule ourselves in some time and make - // another attempt with the already read data but pause reading until - // success - // In this approach we delegate the retry and the backpressure onto the - // caller. This might be enough for most cases. + try { + // busyLooping retries FAIL_NON_SERIALIZED under concurrent senders + this.outboundSink.emitNext(message, Sinks.EmitFailureHandler.busyLooping(Duration.ofMillis(100))); return Mono.empty(); } - else { - return Mono.error(new RuntimeException("Failed to enqueue message")); + catch (Sinks.EmissionException e) { + return Mono.error(new RuntimeException("Failed to enqueue message", e)); } } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java index cc51233..8370e41 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java @@ -10,13 +10,13 @@ import java.io.InputStream; import java.io.InputStreamReader; import java.io.OutputStream; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import io.modelcontextprotocol.json.TypeRef; -import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage; import io.modelcontextprotocol.spec.McpServerSession; @@ -175,11 +175,12 @@ public class StdioServerTransportProvider implements McpServerTransportProvider public Mono sendMessage(McpSchema.JSONRPCMessage message) { return Mono.zip(inboundReady.asMono(), outboundReady.asMono()).then(Mono.defer(() -> { - if (outboundSink.tryEmitNext(message).isSuccess()) { + try { + outboundSink.emitNext(message, Sinks.EmitFailureHandler.busyLooping(Duration.ofMillis(100))); return Mono.empty(); } - else { - return Mono.error(new RuntimeException("Failed to enqueue message")); + catch (Sinks.EmissionException e) { + return Mono.error(new RuntimeException("Failed to enqueue message", e)); } })); } diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java index ab10689..72d0dfc 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java @@ -7,8 +7,11 @@ package io.modelcontextprotocol.server.transport; import java.io.BufferedReader; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.FilterOutputStream; +import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; +import java.io.OutputStream; import java.io.PrintStream; import java.nio.charset.StandardCharsets; import java.time.Duration; @@ -25,7 +28,9 @@ import org.awaitility.Awaitility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import reactor.test.StepVerifier; import static org.assertj.core.api.Assertions.assertThat; @@ -99,7 +104,7 @@ class StdioServerTransportProviderTests { } @Test - void shouldHandleIncomingMessages() throws Exception { + void shouldHandleIncomingMessages() { String jsonMessage = "{\"jsonrpc\":\"2.0\",\"method\":\"test\",\"params\":{},\"id\":1}\n"; InputStream stream = new ByteArrayInputStream(jsonMessage.getBytes(StandardCharsets.UTF_8)); @@ -229,7 +234,7 @@ class StdioServerTransportProviderTests { } @Test - void shouldHandleInvalidJsonMessage() throws Exception { + void shouldHandleInvalidJsonMessage() { // Write an invalid JSON message to the input stream String jsonMessage = "{invalid json}\n"; @@ -248,7 +253,7 @@ class StdioServerTransportProviderTests { } @Test - void shouldRejectInboundMessageExceedingMaxSize() throws Exception { + void shouldRejectInboundMessageExceedingMaxSize() { // A line larger than the configured limit that never terminates with a newline. // BufferedReader#readLine would buffer the whole thing; the bounded reader must // abort instead. @@ -291,7 +296,7 @@ class StdioServerTransportProviderTests { } @Test - void shouldHandleSessionClose() throws Exception { + void shouldHandleSessionClose() { // Set session factory transportProvider.setSessionFactory(sessionFactory); @@ -302,4 +307,69 @@ class StdioServerTransportProviderTests { verify(mockSession).closeGracefully(); } + @Test + void shouldHandleConcurrentSendMessage() throws Exception { + int messageCount = 500; + ByteArrayOutputStream output = new ByteArrayOutputStream(); + CountDownLatch writtenMessages = new CountDownLatch(messageCount); + + // Redirect the transport output to a buffer so we can verify every message lands. + // Writes happen serially on the outbound scheduler, so count the + // newline delimiters as they are written. + OutputStream countingOutput = new FilterOutputStream(output) { + + @Override + public void write(int b) throws IOException { + this.out.write(b); + if (b == '\n') { + writtenMessages.countDown(); + } + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + this.out.write(b, off, len); + for (int i = off; i < off + len; i++) { + if (b[i] == '\n') { + writtenMessages.countDown(); + } + } + } + }; + transportProvider = new StdioServerTransportProvider(McpJsonDefaults.getMapper(), System.in, countingOutput); + + // Capture the inner McpServerTransport handed to the session factory + AtomicReference transportRef = new AtomicReference<>(); + McpServerSession.Factory capturingFactory = transport -> { + transportRef.set(transport); + return mockSession; + }; + + transportProvider.setSessionFactory(capturingFactory); + + McpServerTransport transport = transportRef.get(); + assertThat(transport).isNotNull(); + + // Fan sendMessage out across 16 parallel rails to race against the unicast sink + Flux concurrentSends = Flux.range(0, messageCount) + .parallel(16) + .runOn(Schedulers.parallel()) + .flatMap(i -> transport + .sendMessage( + new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, "test/notification", Map.of())) + .thenReturn(i)) + .sequential(); + + // Every send should complete successfully (no FAIL_NON_SERIALIZED errors) + StepVerifier.create(concurrentSends).expectNextCount(messageCount).verifyComplete(); + + // Wait until the outbound scheduler has actually written all of them + assertThat(writtenMessages.await(30, TimeUnit.SECONDS)) + .as("all %d messages written, %d still missing", messageCount, writtenMessages.getCount()) + .isTrue(); + + // Every message was written as its own newline-delimited JSON line + assertThat(output.toString(StandardCharsets.UTF_8).lines().count()).isEqualTo(messageCount); + } + }