From b563e3176a985d25148f021e107f8c2362cbbc8a Mon Sep 17 00:00:00 2001 From: Pratham Goyal Date: Sat, 24 Oct 2026 08:02:53 +0000 Subject: [PATCH] fix: allow StdioClientTransport to reconnect without a new process The client creates a new session, which connects the transport again, whenever it re-initializes, for instance after an initialize request timed out. StdioClientTransport started a second server process on each connect and subscribed its single-subscriber sinks again, which failed. The first process was left running, and replies still went to the abandoned session, so the client could never initialize. connect now only replaces the inbound message handler once the server process is running. --- .../transport/StdioClientTransport.java | 23 ++++++- .../FirstInitializeIgnoringStdioServer.java | 50 ++++++++++++++ ...oClientTransportReinitializationTests.java | 69 +++++++++++++++++++ 3 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/FirstInitializeIgnoringStdioServer.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportReinitializationTests.java 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 60a8d7c13..ee1072b5c 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 @@ -14,6 +14,8 @@ import java.util.List; import java.util.Set; import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.IntStream; @@ -82,6 +84,12 @@ public class StdioClientTransport implements McpClientTransport { private volatile boolean isClosing = false; + /** Whether the server process has been started by {@link #connect(Function)} */ + private final AtomicBoolean started = new AtomicBoolean(false); + + /** Handler of the messages received from the server process */ + private final AtomicReference, Mono>> inboundMessageHandler = new AtomicReference<>(); + // visible for tests private Consumer stdErrorHandler = error -> logger.info("STDERR Message received: {}", error); @@ -134,8 +142,17 @@ public StdioClientTransport(ServerParameters params, McpJsonMapper jsonMapper, i @Override public Mono connect(Function, Mono> handler) { return Mono.fromRunnable(() -> { + // A client session connects again when it re-initializes, for instance after + // a + // failed initialization: hand it the messages of the already running server + // rather than starting another one. + this.inboundMessageHandler.set(handler); + if (!this.started.compareAndSet(false, true)) { + logger.debug("MCP server already started, replacing the inbound message handler"); + return; + } logger.info("MCP server starting."); - handleIncomingMessages(handler); + handleIncomingMessages(); handleIncomingErrors(); // Prepare command and environment @@ -242,10 +259,10 @@ private void startErrorProcessing() { }); } - private void handleIncomingMessages(Function, Mono> inboundMessageHandler) { + private void handleIncomingMessages() { this.inboundSink.asFlux() .flatMap(message -> Mono.just(message) - .transform(inboundMessageHandler) + .transform(this.inboundMessageHandler.get()) .contextWrite(ctx -> ctx.put("observation", "myObservation"))) .subscribe(); } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/FirstInitializeIgnoringStdioServer.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/FirstInitializeIgnoringStdioServer.java new file mode 100644 index 000000000..d34089379 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/FirstInitializeIgnoringStdioServer.java @@ -0,0 +1,50 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.io.BufferedReader; +import java.io.FileDescriptor; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.InputStreamReader; +import java.io.PrintStream; +import java.nio.charset.StandardCharsets; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Minimal stdio MCP server that never answers the first {@code initialize} request, like + * a server that is too slow to start, and answers every later one. + */ +public final class FirstInitializeIgnoringStdioServer { + + private static final Pattern ID = Pattern.compile("\"id\"\\s*:\\s*(\"[^\"]*\"|-?\\d+)"); + + private FirstInitializeIgnoringStdioServer() { + } + + public static void main(String[] args) throws IOException { + BufferedReader in = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8)); + PrintStream out = new PrintStream(new FileOutputStream(FileDescriptor.out), true, StandardCharsets.UTF_8); + boolean ignoredFirstInitialize = false; + String line; + while ((line = in.readLine()) != null) { + if (!line.contains("\"initialize\"")) { + continue; + } + if (!ignoredFirstInitialize) { + ignoredFirstInitialize = true; + continue; + } + Matcher id = ID.matcher(line); + if (id.find()) { + out.println("{\"jsonrpc\":\"2.0\",\"id\":" + id.group(1) + + ",\"result\":{\"protocolVersion\":\"2025-11-25\",\"capabilities\":{}," + + "\"serverInfo\":{\"name\":\"test-server\",\"version\":\"1.0.0\"}}}"); + } + } + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportReinitializationTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportReinitializationTests.java new file mode 100644 index 000000000..8574f2409 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/StdioClientTransportReinitializationTests.java @@ -0,0 +1,69 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.nio.file.Path; +import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; + +import io.modelcontextprotocol.client.McpAsyncClient; +import io.modelcontextprotocol.client.McpClient; +import io.modelcontextprotocol.json.schema.JsonSchemaValidator; +import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import reactor.test.StepVerifier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; + +/** + * Tests that a {@link StdioClientTransport} can initialize again after a failed + * initialization, without starting another server process. + */ +class StdioClientTransportReinitializationTests { + + private McpAsyncClient client; + + @AfterEach + void tearDown() { + if (this.client != null) { + this.client.close(); + } + } + + @Test + void initializeAgainAfterTimeoutReusesTheServerProcess() { + String java = Path.of(System.getProperty("java.home"), "bin", "java").toString(); + ServerParameters params = ServerParameters.builder(java) + .args("-cp", System.getProperty("java.class.path"), FirstInitializeIgnoringStdioServer.class.getName()) + .build(); + + AtomicInteger startedProcesses = new AtomicInteger(); + StdioClientTransport transport = new StdioClientTransport(params, new GsonMcpJsonMapper()) { + @Override + protected ProcessBuilder getProcessBuilder() { + startedProcesses.incrementAndGet(); + return super.getProcessBuilder(); + } + }; + + this.client = McpClient.async(transport) + .jsonSchemaValidator(mock(JsonSchemaValidator.class)) + .initializationTimeout(Duration.ofSeconds(5)) + .build(); + + // The server does not answer the first initialize request + StepVerifier.create(this.client.initialize()).expectError().verify(Duration.ofSeconds(30)); + + // Retrying must reach the same, still running, server + StepVerifier.create(this.client.initialize()) + .assertNext(result -> assertThat(result.serverInfo().name()).isEqualTo("test-server")) + .verifyComplete(); + + assertThat(startedProcesses).hasValue(1); + } + +}