refactor(mcp): improve StdioServerTransport shutdown and error handling

tzolov · tzolov · commit 9b6a0b537bc6 · 2024-12-29T21:40:53.000+01:00
- 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
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
@@ -171,9 +171,18 @@ private void startInboundProcessing() {
 				}
 			}
 			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 @@ private void startOutboundProcessing() {
 							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 Mono<Void> closeGracefully() {
 		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.fromRunnable(() -> {
-				try {
-
-					inboundScheduler.dispose();
-					outboundScheduler.dispose();
+		}).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();
+
+				// Wait for schedulers to terminate
+				if (!inboundScheduler.isDisposed()) {
+					inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+				}
+				if (!outboundScheduler.isDisposed()) {
+					outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+				}
 
-					// Wait for schedulers to terminate
-					if (!inboundScheduler.isDisposed()) {
-						inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+				// Only after schedulers are disposed, close the streams
+				try {
+					if (inputStream != System.in) {
+						inputStream.close();
 					}
-					if (!outboundScheduler.isDisposed()) {
-						outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
+					if (outputStream != System.out) {
+						outputStream.flush();
+						outputStream.close();
 					}
-
-					logger.info("Graceful shutdown completed");
 				}
-				catch (Exception e) {
-					logger.error("Error during graceful shutdown", e);
+				catch (IOException e) {
+					// Log but don't throw since we're shutting down
+					logger.debug("Error closing streams during 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
