Make STDIO transport sendMessage thread-safe (#304)

Concurrent calls to `sendMessage` on the STDIO transports raced on the
unicast sink and sporadically failed with `FAIL_NON_SERIALIZED`, dropping
the message with `RuntimeException("Failed to enqueue message")`.

Switch `tryEmitNext` to `emitNext` with a `busyLooping` EmitFailureHandler
so concurrent emissions are serialized as recommended by the Reactor
maintainers. Applied on both `StdioServerTransportProvider` and
`StdioClientTransport`.

Co-authored-by: Nicolas Quinquenel <nicolas.quinquenel@sonarsource.com>
Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

Author: nquinquenel

mcp-core/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransport.java

@@ -258,16 +258,13 @@ private void handleIncomingErrors() {
 
 	@Override
 	public Mono<Void> 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));
 		}
 	}
 

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java

@@ -10,13 +10,13 @@
 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 StdioMcpSessionTransport() {
 		public Mono<Void> 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));
 				}
 			}));
 		}

mcp-test/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java

@@ -7,8 +7,11 @@
 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.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 @@ void shouldCreateSessionWhenSessionFactoryIsSet() {
 	}
 
 	@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 @@ void shouldHandleNotificationBeforeSessionFactoryIsSet() {
 	}
 
 	@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 @@ void shouldHandleInvalidJsonMessage() throws Exception {
 	}
 
 	@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 @@ void shouldRejectNonPositiveMaxSize() {
 	}
 
 	@Test
-	void shouldHandleSessionClose() throws Exception {
+	void shouldHandleSessionClose() {
 		// Set session factory
 		transportProvider.setSessionFactory(sessionFactory);
 
@@ -302,4 +307,69 @@ void shouldHandleSessionClose() throws Exception {
 		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<McpServerTransport> 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<Integer> 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);
+	}
+
 }