refactor(mcp) Disable stdio server sink draining on close

tzolov · tzolov · commit dfa7435ad888 · 2024-12-30T09:23:42.000+01:00
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
@@ -248,46 +248,50 @@ public Mono<Void> closeGracefully() {
 		return Mono.fromRunnable(() -> {
 			isClosing = true;
 			logger.debug("Initiating graceful shutdown");
-		}).then(Mono.defer(() -> {
-			// First complete the sinks to stop processing
-			inboundSink.tryEmitComplete();
-			outboundSink.tryEmitComplete();
-			return Mono.delay(Duration.ofMillis(200));
-		})).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));
-				}
-
-				// Only after schedulers are disposed, close the streams
+		})
+			// .then(Mono.defer(() -> {
+			// // First complete the sinks to stop processing
+			// inboundSink.tryEmitComplete();
+			// outboundSink.tryEmitComplete();
+			// return Mono.delay(Duration.ofMillis(100));
+			// }))
+			.then(Mono.fromRunnable(() -> {
 				try {
-					if (inputStream != System.in) {
-						inputStream.close();
+					// 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));
+					}
+
+					// Only after schedulers are disposed, close the streams
+					try {
+						if (inputStream != System.in) {
+							inputStream.close();
+						}
+						if (outputStream != System.out) {
+							outputStream.flush();
+							outputStream.close();
+						}
 					}
-					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 (IOException e) {
-					// Log but don't throw since we're shutting down
-					logger.debug("Error closing streams during shutdown", e);
+				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
