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); + } + +}