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 8370e41ed..c702ebd23 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 @@ -174,7 +174,7 @@ public StdioMcpSessionTransport() { @Override public Mono sendMessage(McpSchema.JSONRPCMessage message) { - return Mono.zip(inboundReady.asMono(), outboundReady.asMono()).then(Mono.defer(() -> { + return Mono.when(inboundReady.asMono(), outboundReady.asMono()).then(Mono.defer(() -> { try { outboundSink.emitNext(message, Sinks.EmitFailureHandler.busyLooping(Duration.ofMillis(100))); return Mono.empty(); 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 72d0dfcf8..859ad7c68 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 @@ -13,6 +13,7 @@ import java.io.InputStreamReader; import java.io.OutputStream; import java.io.PrintStream; +import java.lang.reflect.Field; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.Map; @@ -30,6 +31,7 @@ import org.junit.jupiter.api.Test; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; import reactor.core.scheduler.Schedulers; import reactor.test.StepVerifier; @@ -307,6 +309,34 @@ void shouldHandleSessionClose() { verify(mockSession).closeGracefully(); } + @Test + @SuppressWarnings("unchecked") + void sendMessageWaitsForBothInboundAndOutboundReadinessSignals() throws Exception { + transportProvider = new StdioServerTransportProvider(McpJsonDefaults.getMapper(), System.in, + testOutPrintStream); + Class transportClass = Class + .forName(StdioServerTransportProvider.class.getName() + "$StdioMcpSessionTransport"); + var constructor = transportClass.getDeclaredConstructor(StdioServerTransportProvider.class); + constructor.setAccessible(true); + McpServerTransport transport = (McpServerTransport) constructor.newInstance(transportProvider); + + Field inboundReadyField = StdioServerTransportProvider.class.getDeclaredField("inboundReady"); + inboundReadyField.setAccessible(true); + Sinks.One inboundReady = (Sinks.One) inboundReadyField.get(transportProvider); + + Field outboundReadyField = transportClass.getDeclaredField("outboundReady"); + outboundReadyField.setAccessible(true); + Sinks.One outboundReady = (Sinks.One) outboundReadyField.get(transport); + + StepVerifier + .create(transport.sendMessage( + new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, "test/notification", Map.of()))) + .then(() -> inboundReady.tryEmitValue(null)) + .expectNoEvent(Duration.ofMillis(100)) + .then(() -> outboundReady.tryEmitValue(null)) + .verifyComplete(); + } + @Test void shouldHandleConcurrentSendMessage() throws Exception { int messageCount = 500;