From 48ab97bf237795c2f5ba1ee37351fa225d3dafb6 Mon Sep 17 00:00:00 2001
From: Christian Tzolov <christian.tzolov@broadcom.com>
Date: Fri, 24 Jan 2025 13:08:41 +0100
Subject: [PATCH] fix(transport): Add cleanup logic for inbound processing in
 StdioServerTransport

Resolves #72

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
---
 .../transport/StdioServerTransport.java       | 91 ++++++++++---------
 1 file changed, 48 insertions(+), 43 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 5f8e78b..e9b88d2 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
@@ -118,6 +118,11 @@ public class StdioServerTransport implements ServerMcpTransport {
 			.flatMap(message -> Mono.just(message)
 				.transform(inboundMessageHandler)
 				.contextWrite(ctx -> ctx.put("observation", "myObservation")))
+			.doOnComplete(() -> {
+				this.outboundSink.tryEmitComplete();
+				this.inboundScheduler.dispose();
+				this.outboundScheduler.dispose();
+			})
 			.subscribe();
 	}
 
@@ -184,51 +189,51 @@ public class StdioServerTransport implements ServerMcpTransport {
 	 */
 	private void startOutboundProcessing() {
 		Function<Flux<JSONRPCMessage>, Flux<JSONRPCMessage>> outboundConsumer = messages -> messages // @formatter:off
-			.doOnSubscribe(subscription -> outboundReady.tryEmitValue(null))
-			.publishOn(outboundScheduler)
-			.handle((message, sink) -> {
-				if (message != null && !isClosing) {
-					try {
-						String jsonMessage = objectMapper.writeValueAsString(message);
-						// Escape any embedded newlines in the JSON message as per spec
-						jsonMessage = jsonMessage.replace("\r\n", "\\n").replace("\n", "\\n").replace("\r", "\\n");
+			 .doOnSubscribe(subscription -> outboundReady.tryEmitValue(null))
+			 .publishOn(outboundScheduler)
+			 .handle((message, sink) -> {
+				 if (message != null && !isClosing) {
+					 try {
+						 String jsonMessage = objectMapper.writeValueAsString(message);
+						 // Escape any embedded newlines in the JSON message as per spec
+						 jsonMessage = jsonMessage.replace("\r\n", "\\n").replace("\n", "\\n").replace("\r", "\\n");
 
-						synchronized (outputStream) {
-							outputStream.write(jsonMessage.getBytes(StandardCharsets.UTF_8));
-							outputStream.write("\n".getBytes(StandardCharsets.UTF_8));
-							outputStream.flush();
-						}
-						sink.next(message);
-					}
-					catch (IOException e) {
-						if (!isClosing) {
-							logger.error("Error writing message", e);
-							sink.error(new RuntimeException(e));
-						}
-						else {
-							logger.debug("Stream closed during shutdown", e);
-						}
-					}
-				}
-				else if (isClosing) {
-					sink.complete();
-				}
-			})
-			.doOnComplete(() -> {
-				isClosing = true;
-				outboundSink.tryEmitComplete();
-			})
-			.doOnError(e -> {
-				if (!isClosing) {
-					logger.error("Error in outbound processing", e);
-					isClosing = true;
-					outboundSink.tryEmitComplete();
-				}
-			})
-			.map(msg -> (JSONRPCMessage) msg);
+						 synchronized (outputStream) {
+							 outputStream.write(jsonMessage.getBytes(StandardCharsets.UTF_8));
+							 outputStream.write("\n".getBytes(StandardCharsets.UTF_8));
+							 outputStream.flush();
+						 }
+						 sink.next(message);
+					 }
+					 catch (IOException e) {
+						 if (!isClosing) {
+							 logger.error("Error writing message", e);
+							 sink.error(new RuntimeException(e));
+						 }
+						 else {
+							 logger.debug("Stream closed during shutdown", e);
+						 }
+					 }
+				 }
+				 else if (isClosing) {
+					 sink.complete();
+				 }
+			 })
+			 .doOnComplete(() -> {
+				 isClosing = true;
+				 outboundSink.tryEmitComplete();
+			 })
+			 .doOnError(e -> {
+				 if (!isClosing) {
+					 logger.error("Error in outbound processing", e);
+					 isClosing = true;
+					 outboundSink.tryEmitComplete();
+				 }
+			 })
+			 .map(msg -> (JSONRPCMessage) msg);
 
-			outboundConsumer.apply(outboundSink.asFlux()).subscribe();
-	} // @formatter:on
+			 outboundConsumer.apply(outboundSink.asFlux()).subscribe();
+	 } // @formatter:on
 
 	@Override
 	public Mono<Void> closeGracefully() {