From 9b6a0b537bc6a72b54ae1433daef70fa2e817c80 Mon Sep 17 00:00:00 2001
From: Christian Tzolov <ctzolov@vmware.com>
Date: Sun, 29 Dec 2024 21:40:51 +0100
Subject: [PATCH] refactor(mcp): improve StdioServerTransport shutdown and
 error handling

- Add specific error categorization for pipe closed errors
- Enhance shutdown sequence with proper stream and scheduler cleanup
- Add debug logging for shutdown-related stream errors
- Improve isClosing state management
---
 .../transport/StdioServerTransport.java       | 73 ++++++++++++-------
 1 file changed, 48 insertions(+), 25 deletions(-)

diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java b/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java
index 4d9c979..dff9727 100644
--- a/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java
+++ b/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java
@@ -171,9 +171,18 @@ public class StdioServerTransport implements McpTransport {
 				}
 			}
 			catch (IOException e) {
-				if (!isClosing) {
+				// Check isClosing before the error occurs to properly categorize it
+				boolean wasClosing = isClosing;
+				isClosing = true;
+				if (!wasClosing && e.getMessage().equals("Pipe closed")) {
+					logger.debug("Stream closed during shutdown", e);
+				}
+				else if (!wasClosing) {
 					logger.error("Error reading from stdin", e);
 				}
+				else {
+					logger.debug("Stream error during shutdown", e);
+				}
 			}
 			finally {
 				isClosing = true;
@@ -209,6 +218,9 @@ public class StdioServerTransport implements McpTransport {
 							logger.error("Error writing message", e);
 							sink.error(new RuntimeException(e));
 						}
+						else {
+							logger.debug("Stream closed during shutdown", e);
+						}
 					}
 				}
 				else if (isClosing) {
@@ -236,35 +248,46 @@ public class StdioServerTransport implements McpTransport {
 		return Mono.fromRunnable(() -> {
 			isClosing = true;
 			logger.debug("Initiating graceful shutdown");
-		})
-			// .then(Mono.defer(() -> {
-			// inboundSink.tryEmitComplete();
-			// outboundSink.tryEmitComplete();
-			// return Mono.delay(Duration.ofMillis(100));
-			// }))
+		}).then(Mono.defer(() -> {
+			// First complete the sinks to stop processing
+			inboundSink.tryEmitComplete();
+			outboundSink.tryEmitComplete();
+			return Mono.delay(Duration.ofMillis(100));
+		})).then(Mono.fromRunnable(() -> {
+			try {
+				// Dispose schedulers first
+				inboundScheduler.dispose();
+				outboundScheduler.dispose();
 
-			.then(Mono.fromRunnable(() -> {
+				// Wait for schedulers to terminate
+				if (!inboundScheduler.isDisposed()) {
+					inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+				}
+				if (!outboundScheduler.isDisposed()) {
+					outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+				}
+
+				// Only after schedulers are disposed, close the streams
 				try {
-
-					inboundScheduler.dispose();
-					outboundScheduler.dispose();
-
-					// Wait for schedulers to terminate
-					if (!inboundScheduler.isDisposed()) {
-						inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+					if (inputStream != System.in) {
+						inputStream.close();
 					}
-					if (!outboundScheduler.isDisposed()) {
-						outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+					if (outputStream != System.out) {
+						outputStream.flush();
+						outputStream.close();
 					}
+				}
+				catch (IOException e) {
+					// Log but don't throw since we're shutting down
+					logger.debug("Error closing streams during shutdown", e);
+				}
 
-					logger.info("Graceful shutdown completed");
-				}
-				catch (Exception e) {
-					logger.error("Error during graceful shutdown", e);
-				}
-			}))
-			.then()
-			.subscribeOn(Schedulers.boundedElastic());
+				logger.info("Graceful shutdown completed");
+			}
+			catch (Exception e) {
+				logger.error("Error during graceful shutdown", e);
+			}
+		})).then().subscribeOn(Schedulers.boundedElastic());
 	}
 
 	@Override