fix(mcp): Improve transport shutdown and error handling

- Add isClosing flag to coordinate graceful shutdown
- Synchronize output stream writes to prevent race conditions
- Enhance error handling in stream processing
- Add delay for pending message processing during shutdown
- Improve connection state tracking in tests

Author: tzolov

mcp/src/main/java/org/springframework/ai/mcp/client/transport/StdioClientTransport.java

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -77,6 +78,8 @@ public class StdioClientTransport implements McpTransport {
 
 	private final Sinks.Many<String> errorSink;
 
+	private volatile boolean isClosing = false;
+
 	// visible for tests
 	private Consumer<String> errorHandler = error -> logger.error("Error received: {}", error);
 
@@ -200,19 +203,32 @@ private void startErrorProcessing() {
 			try (BufferedReader processErrorReader = new BufferedReader(
 					new InputStreamReader(process.getErrorStream()))) {
 				String line;
-				while ((line = processErrorReader.readLine()) != null) {
+				while (!isClosing && (line = processErrorReader.readLine()) != null) {
 					try {
 						logger.error("Received error line: {}", line);
-						// TODO: handle errors, etc.
-						this.errorSink.tryEmitNext(line);
+						if (!this.errorSink.tryEmitNext(line).isSuccess()) {
+							if (!isClosing) {
+								logger.error("Failed to emit error message");
+							}
+							break;
+						}
 					}
 					catch (Exception e) {
-						throw new RuntimeException(e);
+						if (!isClosing) {
+							logger.error("Error processing error message", e);
+						}
+						break;
 					}
 				}
 			}
 			catch (IOException e) {
-				throw new RuntimeException(e);
+				if (!isClosing) {
+					logger.error("Error reading from error stream", e);
+				}
+			}
+			finally {
+				isClosing = true;
+				errorSink.tryEmitComplete();
 			}
 		});
 	}
@@ -254,21 +270,32 @@ private void startInboundProcessing() {
 		this.inboundScheduler.schedule(() -> {
 			try (BufferedReader processReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
 				String line;
-				while ((line = processReader.readLine()) != null) {
+				while (!isClosing && (line = processReader.readLine()) != null) {
 					try {
 						JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
 						if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
-							// TODO: Back off, reschedule, give up?
-							throw new RuntimeException("Failed to enqueue message");
+							if (!isClosing) {
+								logger.error("Failed to enqueue inbound message");
+							}
+							break;
 						}
 					}
 					catch (Exception e) {
-						throw new RuntimeException(e);
+						if (!isClosing) {
+							logger.error("Error processing inbound message", e);
+						}
+						break;
 					}
 				}
 			}
 			catch (IOException e) {
-				throw new RuntimeException(e);
+				if (!isClosing) {
+					logger.error("Error reading from input stream", e);
+				}
+			}
+			finally {
+				isClosing = true;
+				inboundSink.tryEmitComplete();
 			}
 		});
 	}
@@ -284,7 +311,7 @@ private void startOutboundProcessing() {
 			// want to ensure that the actual writing happens on a dedicated thread
 			.publishOn(outboundScheduler)
 			.handle((message, s) -> {
-				if (message != null) {
+				if (message != null && !isClosing) {
 					try {
 						String jsonMessage = objectMapper.writeValueAsString(message);
 						// Escape any embedded newlines in the JSON message as per spec:
@@ -293,9 +320,12 @@ private void startOutboundProcessing() {
 						// embedded newlines.
 						jsonMessage = jsonMessage.replace("\r\n", "\\n").replace("\n", "\\n").replace("\r", "\\n");
 
-						this.process.getOutputStream().write(jsonMessage.getBytes(StandardCharsets.UTF_8));
-						this.process.getOutputStream().write("\n".getBytes(StandardCharsets.UTF_8));
-						this.process.getOutputStream().flush();
+						var os = this.process.getOutputStream();
+						synchronized (os) {
+							os.write(jsonMessage.getBytes(StandardCharsets.UTF_8));
+							os.write("\n".getBytes(StandardCharsets.UTF_8));
+							os.flush();
+						}
 						s.next(message);
 					}
 					catch (IOException e) {
@@ -306,7 +336,16 @@ private void startOutboundProcessing() {
 	}
 
 	protected void handleOutbound(Function<Flux<JSONRPCMessage>, Flux<JSONRPCMessage>> outboundConsumer) {
-		outboundConsumer.apply(outboundSink.asFlux()).subscribe();
+		outboundConsumer.apply(outboundSink.asFlux()).doOnComplete(() -> {
+			isClosing = true;
+			outboundSink.tryEmitComplete();
+		}).doOnError(e -> {
+			if (!isClosing) {
+				logger.error("Error in outbound processing", e);
+				isClosing = true;
+				outboundSink.tryEmitComplete();
+			}
+		}).subscribe();
 	}
 
 	/**
@@ -317,7 +356,18 @@ protected void handleOutbound(Function<Flux<JSONRPCMessage>, Flux<JSONRPCMessage
 	 */
 	@Override
 	public Mono<Void> closeGracefully() {
-		return Mono.fromFuture(() -> {
+		return Mono.fromRunnable(() -> {
+			isClosing = true;
+			logger.debug("Initiating graceful shutdown");
+		}).then(Mono.defer(() -> {
+			// First complete all sinks to stop accepting new messages
+			inboundSink.tryEmitComplete();
+			outboundSink.tryEmitComplete();
+			errorSink.tryEmitComplete();
+
+			// Give a short time for any pending messages to be processed
+			return Mono.delay(Duration.ofMillis(100));
+		})).then(Mono.fromFuture(() -> {
 			logger.info("Sending TERM to process");
 			if (this.process != null) {
 				this.process.destroy();
@@ -326,16 +376,23 @@ public Mono<Void> closeGracefully() {
 			else {
 				return CompletableFuture.failedFuture(new RuntimeException("Process not started"));
 			}
-		}).doOnNext(process -> {
+		})).doOnNext(process -> {
 			if (process.exitValue() != 0) {
 				logger.warn("Process terminated with code " + process.exitValue());
 			}
 		}).then(Mono.fromRunnable(() -> {
-			// The Threads are blocked on readLine so disposeGracefully would not
-			// interrupt them, therefore we issue an async hard dispose.
-			inboundScheduler.dispose();
-			errorScheduler.dispose();
-			outboundScheduler.dispose();
+			try {
+				// The Threads are blocked on readLine so disposeGracefully would not
+				// interrupt them, therefore we issue an async hard dispose.
+				inboundScheduler.dispose();
+				errorScheduler.dispose();
+				outboundScheduler.dispose();
+
+				logger.info("Graceful shutdown completed");
+			}
+			catch (Exception e) {
+				logger.error("Error during graceful shutdown", e);
+			}
 		})).then().subscribeOn(Schedulers.boundedElastic());
 	}
 

mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java

@@ -176,9 +176,7 @@ private void startInboundProcessing() {
 				}
 			}
 			finally {
-				if (!isClosing) {
-					isClosing = true;
-				}
+				isClosing = true;
 				inboundSink.tryEmitComplete();
 			}
 		});
@@ -218,9 +216,7 @@ else if (isClosing) {
 				}
 			})
 			.doOnComplete(() -> {
-				if (!isClosing) {
-					isClosing = true;
-				}
+				isClosing = true;
 				outboundSink.tryEmitComplete();
 			})
 			.doOnError(e -> {
@@ -240,29 +236,35 @@ public Mono<Void> closeGracefully() {
 		return Mono.fromRunnable(() -> {
 			isClosing = true;
 			logger.debug("Initiating graceful shutdown");
-			// }).then(Mono.delay(Duration.ofMillis(100))).then(Mono.fromRunnable(() -> {
-		}).then(Mono.fromRunnable(() -> {
-			try {
-				// inboundSink.tryEmitComplete();
-				// outboundSink.tryEmitComplete();
-
-				inboundScheduler.dispose();
-				outboundScheduler.dispose();
-
-				// Wait for schedulers to terminate
-				if (!inboundScheduler.isDisposed()) {
-					inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+		})
+			// .then(Mono.defer(() -> {
+			// inboundSink.tryEmitComplete();
+			// outboundSink.tryEmitComplete();
+			// return Mono.delay(Duration.ofMillis(100));
+			// }))
+
+			.then(Mono.fromRunnable(() -> {
+				try {
+
+					inboundScheduler.dispose();
+					outboundScheduler.dispose();
+
+					// Wait for schedulers to terminate
+					if (!inboundScheduler.isDisposed()) {
+						inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+					}
+					if (!outboundScheduler.isDisposed()) {
+						outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+					}
+
+					logger.info("Graceful shutdown completed");
 				}
-				if (!outboundScheduler.isDisposed()) {
-					outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+				catch (Exception e) {
+					logger.error("Error during graceful shutdown", e);
 				}
-
-				logger.info("Graceful shutdown completed");
-			}
-			catch (Exception e) {
-				logger.error("Error during graceful shutdown", e);
-			}
-		})).then().subscribeOn(Schedulers.boundedElastic());
+			}))
+			.then()
+			.subscribeOn(Schedulers.boundedElastic());
 	}
 
 	@Override

mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java

@@ -49,11 +49,11 @@ private static class MockMcpTransport implements McpTransport {
 
 		private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
 
-		private Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
+		private final Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
 
-		private Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
+		private final Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
 
-		private Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
+		private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
 
 		public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
 			if (inbound.tryEmitNext(message).isFailure()) {
@@ -82,17 +82,29 @@ public McpSchema.JSONRPCMessage getLastSentMessage() {
 			return outboundView.blockFirst();
 		}
 
+		private volatile boolean connected = false;
+
 		@Override
 		public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
+			if (connected) {
+				return Mono.error(new IllegalStateException("Already connected"));
+			}
+			connected = true;
 			return inbound.asFlux()
 				.publishOn(Schedulers.boundedElastic())
 				.flatMap(message -> Mono.just(message).transform(handler))
+				.doFinally(signal -> connected = false)
 				.then();
 		}
 
 		@Override
 		public Mono<Void> closeGracefully() {
-			return Mono.empty();
+			return Mono.defer(() -> {
+				connected = false;
+				outgoing.tryEmitComplete();
+				inbound.tryEmitComplete();
+				return Mono.empty();
+			});
 		}
 
 		@Override