fix listening sse stream close blocking

Author: taobaorun

mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java

@@ -12,6 +12,7 @@
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpStreamableServerSession;
+import io.modelcontextprotocol.spec.McpStreamableServerSession.McpStreamableServerSessionStream;
 import io.modelcontextprotocol.spec.McpStreamableServerTransport;
 import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
 import io.modelcontextprotocol.spec.ProtocolVersions;
@@ -199,9 +200,16 @@ private Mono<ServerResponse> handleGet(ServerRequest request) {
 				.body(Flux.<ServerSentEvent<?>>create(sink -> {
 					WebFluxStreamableMcpSessionTransport sessionTransport = new WebFluxStreamableMcpSessionTransport(
 							sink);
-					McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
-						.listeningStream(sessionTransport);
-					sink.onDispose(listeningStream::close);
+					session.listeningStream(sessionTransport)
+						.doOnNext(serverSessionStream -> sink
+							.onDispose(() -> serverSessionStream.closeGracefully().subscribe(v -> {
+							}, error -> logger.warn("Failed to close listening stream gracefully", error))))
+						.doOnError(error -> {
+							logger.error("Failed to create listening stream", error);
+							sink.error(error);
+						})
+						.subscribe(serverSessionStream -> logger.debug("Listening stream created successfully"),
+								sink::error);
 				}), ServerSentEvent.class);
 
 		}).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext));

mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java

@@ -10,6 +10,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.ReentrantLock;
 
+import io.modelcontextprotocol.spec.McpStreamableServerSession.McpStreamableServerSessionStream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpStatus;
@@ -289,13 +290,15 @@ private ServerResponse handleGet(ServerRequest request) {
 				}
 				else {
 					// Establish new listening stream
-					McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
+					Mono<McpStreamableServerSession.McpStreamableServerSessionStream> listeningStream = session
 						.listeningStream(sessionTransport);
-
-					sseBuilder.onComplete(() -> {
+					listeningStream.subscribe(serverSessionStream -> sseBuilder.onComplete(() -> {
 						logger.debug("SSE connection completed for session: {}", sessionId);
-						listeningStream.close();
-					});
+						serverSessionStream.close();
+					}), error -> {
+						sseBuilder.error(error);
+						logger.error("Failed to create listening stream", error);
+					}, () -> logger.debug("Listening stream created successfully"));
 				}
 			}, Duration.ZERO);
 		}

mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java

@@ -13,6 +13,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.ReentrantLock;
 
+import io.modelcontextprotocol.spec.McpStreamableServerSession.McpStreamableServerSessionStream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -316,33 +317,36 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
 			}
 			else {
 				// Establish new listening stream
-				McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
-					.listeningStream(sessionTransport);
-
-				asyncContext.addListener(new jakarta.servlet.AsyncListener() {
-					@Override
-					public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException {
-						logger.debug("SSE connection completed for session: {}", sessionId);
-						listeningStream.close();
-					}
-
-					@Override
-					public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException {
-						logger.debug("SSE connection timed out for session: {}", sessionId);
-						listeningStream.close();
-					}
+				session.listeningStream(sessionTransport)
+					.doOnNext(serverSessionStream -> asyncContext.addListener(new jakarta.servlet.AsyncListener() {
+						@Override
+						public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException {
+							logger.debug("SSE connection completed for session: {}", sessionId);
+							serverSessionStream.close();
+						}
+
+						@Override
+						public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException {
+							logger.debug("SSE connection timed out for session: {}", sessionId);
+							serverSessionStream.close();
+						}
+
+						@Override
+						public void onError(jakarta.servlet.AsyncEvent event) throws IOException {
+							logger.debug("SSE connection error for session: {}", sessionId);
+							serverSessionStream.close();
+						}
+
+						@Override
+						public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException {
+							// No action needed
+						}
+					}))
+					.doOnError(error -> {
+						logger.error("Failed to create listening stream", error);
+					})
+					.subscribe(serverSessionStream -> logger.debug("Listening stream created successfully"));
 
-					@Override
-					public void onError(jakarta.servlet.AsyncEvent event) throws IOException {
-						logger.debug("SSE connection error for session: {}", sessionId);
-						listeningStream.close();
-					}
-
-					@Override
-					public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException {
-						// No action needed
-					}
-				});
 			}
 		}
 		catch (Exception e) {

mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java

@@ -134,16 +134,16 @@ public Mono<Void> delete() {
 	 * @param transport The dedicated SSE transport stream
 	 * @return a stream representation
 	 */
-	public McpStreamableServerSessionStream listeningStream(McpStreamableServerTransport transport) {
+	public Mono<McpStreamableServerSessionStream> listeningStream(McpStreamableServerTransport transport) {
 		McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport);
-		McpLoggableSession listenedStream = this.listeningStreamRef.getAndSet(listeningStream);
-		if (listenedStream != null) {
+		McpLoggableSession oldStream = this.listeningStreamRef.getAndSet(listeningStream);
+		if (oldStream != null) {
 			logger.debug(
 					"Listening stream already exists for this session:{} and will be closed to make way for the new listening SSE stream",
 					this.id);
-			listenedStream.closeGracefully().block();
+			return oldStream.closeGracefully().thenReturn(listeningStream);
 		}
-		return listeningStream;
+		return Mono.just(listeningStream);
 	}
 
 	// TODO: keep track of history by keeping a map from eventId to stream and then