From 7fb12e2513b3f0d4138b9de16553417ad6236371 Mon Sep 17 00:00:00 2001
From: liruixin <96371341+liruixin@users.noreply.github.com>
Date: Wed, 12 Nov 2025 10:46:33 +0800
Subject: [PATCH] Fix: Added AsyncListener to handle completion, timeout, and
 error events for AsyncContext.

---
 ...HttpServletSseServerTransportProvider.java | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java
index 4739e231a..ab0c5d0c1 100644
--- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java
+++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java
@@ -27,6 +27,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;
@@ -232,6 +234,40 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
 		AsyncContext asyncContext = request.startAsync();
 		asyncContext.setTimeout(0);
 
+		asyncContext.addListener(new AsyncListener() {
+            @Override
+            public void onComplete(AsyncEvent event) {
+                logger.debug("AsyncContext completed for session {}", sessionId);
+                sessions.remove(sessionId);
+            }
+
+            @Override
+            public void onTimeout(AsyncEvent event) {
+                logger.warn("Session {} timeout, cleaning up", sessionId);
+                sessions.remove(sessionId);
+                try {
+                    event.getAsyncContext().complete();
+                } catch (Exception e) {
+                    logger.error("Error completing async context on timeout", e);
+                }
+            }
+
+            @Override
+            public void onError(AsyncEvent event) {
+                logger.error("AsyncContext error for session {}: {}", sessionId, event.getThrowable().getMessage());
+                sessions.remove(sessionId);
+                try {
+                    event.getAsyncContext().complete();
+                } catch (Exception e) {
+                    logger.error("Error completing async context on error", e);
+                }
+            }
+
+            @Override
+            public void onStartAsync(AsyncEvent event) {
+            }
+        });
+
 		PrintWriter writer = response.getWriter();
 
 		// Create a new session transport