Close servlet streamable HTTP transports on async lifecycle events

Author: lxq19991111

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

@@ -28,6 +28,8 @@
 import io.modelcontextprotocol.util.Assert;
 import io.modelcontextprotocol.util.KeepAliveScheduler;
 import jakarta.servlet.AsyncContext;
+import jakarta.servlet.AsyncEvent;
+import jakarta.servlet.AsyncListener;
 import jakarta.servlet.ServletException;
 import jakarta.servlet.annotation.WebServlet;
 import jakarta.servlet.http.HttpServlet;
@@ -329,6 +331,7 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
 			// Check if this is a replay request
 			if (request.getHeader(HttpHeaders.LAST_EVENT_ID) != null) {
 				String lastId = request.getHeader(HttpHeaders.LAST_EVENT_ID);
+				registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close);
 
 				try {
 					session.replay(lastId)
@@ -342,44 +345,21 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
 							}
 							catch (Exception e) {
 								logger.error("Failed to replay message: {}", e.getMessage());
-								asyncContext.complete();
+								sessionTransport.close();
 							}
 						});
 				}
 				catch (Exception e) {
 					logger.error("Failed to replay messages: {}", e.getMessage());
-					asyncContext.complete();
+					sessionTransport.close();
 				}
 			}
 			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();
-					}
-
-					@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
-					}
-				});
+				registerAsyncLifecycle(asyncContext, sessionId, listeningStream::close);
 			}
 		}
 		catch (Exception e) {
@@ -530,6 +510,7 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
 
 				HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
 						sessionId, asyncContext, response.getWriter());
+				registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close);
 
 				try {
 					session.responseStream(jsonrpcRequest, sessionTransport)
@@ -538,7 +519,7 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
 				}
 				catch (Exception e) {
 					logger.error("Failed to handle request stream: {}", e.getMessage());
-					asyncContext.complete();
+					sessionTransport.close();
 				}
 			}
 			else {
@@ -571,6 +552,32 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
 		}
 	}
 
+	private void registerAsyncLifecycle(AsyncContext asyncContext, String sessionId, Runnable onClose) {
+		asyncContext.addListener(new AsyncListener() {
+			@Override
+			public void onComplete(AsyncEvent event) throws IOException {
+				logger.debug("SSE async context completed for session: {}", sessionId);
+				onClose.run();
+			}
+
+			@Override
+			public void onTimeout(AsyncEvent event) throws IOException {
+				logger.debug("SSE async context timed out for session: {}", sessionId);
+				onClose.run();
+			}
+
+			@Override
+			public void onError(AsyncEvent event) throws IOException {
+				logger.debug("SSE async context errored for session: {}", sessionId);
+				onClose.run();
+			}
+
+			@Override
+			public void onStartAsync(AsyncEvent event) throws IOException {
+			}
+		});
+	}
+
 	/**
 	 * Handles DELETE requests for session deletion.
 	 * @param request The HTTP servlet request
@@ -761,8 +768,7 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId
 				}
 				catch (Exception e) {
 					logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage());
-					HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId);
-					this.asyncContext.complete();
+					this.close();
 				}
 				finally {
 					lock.unlock();
@@ -806,8 +812,6 @@ public void close() {
 				}
 
 				this.closed = true;
-
-				// HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId);
 				this.asyncContext.complete();
 				logger.debug("Successfully completed async context for session {}", sessionId);
 			}