Fix: SseServerTransport.handle_sse never ends and can lead to memory leaks

Author: PWZER

examples/servers/simple-tool/mcp_simple_tool/server.py

@@ -60,17 +60,24 @@ async def list_tools() -> list[types.Tool]:
     if transport == "sse":
         from mcp.server.sse import SseServerTransport
         from starlette.applications import Starlette
+        from starlette.requests import Request
+        from starlette.responses import Response
         from starlette.routing import Mount, Route
 
         sse = SseServerTransport("/messages/")
 
-        async def handle_sse(request):
-            async with sse.connect_sse(
-                request.scope, request.receive, request._send
-            ) as streams:
-                await app.run(
-                    streams[0], streams[1], app.create_initialization_options()
-                )
+        async def handle_sse(request: Request):
+            with anyio.CancelScope() as cancel_scope:
+                async with sse.connect_sse(
+                    request.scope,
+                    request.receive,
+                    request._send,
+                    lambda: cancel_scope.cancel(),
+                ) as streams:
+                    await app.run(
+                        streams[0], streams[1], app.create_initialization_options()
+                    )
+            return Response(status_code=200)
 
         starlette_app = Starlette(
             debug=True,

src/mcp/server/sse.py

@@ -32,6 +32,7 @@ async def handle_sse(request):
 """
 
 import logging
+from collections.abc import Callable
 from contextlib import asynccontextmanager
 from typing import Any
 from urllib.parse import quote
@@ -41,6 +42,7 @@ async def handle_sse(request):
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import ValidationError
 from sse_starlette import EventSourceResponse
+from starlette.background import BackgroundTask
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.types import Receive, Scope, Send
@@ -79,7 +81,13 @@ def __init__(self, endpoint: str) -> None:
         logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
 
     @asynccontextmanager
-    async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
+    async def connect_sse(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+        callback: Callable[[], None] | None = None,
+    ):
         if scope["type"] != "http":
             logger.error("connect_sse received non-HTTP request")
             raise ValueError("connect_sse can only handle HTTP requests")
@@ -120,9 +128,19 @@ async def sse_writer():
                         }
                     )
 
+        async def _remove_stream_writer() -> None:
+            await read_stream_writer.aclose()
+            await write_stream_reader.aclose()
+            del self._read_stream_writers[session_id]
+            if callback:
+                callback()
+            logger.debug(f"Closed SSE session with ID: {session_id}")
+
         async with anyio.create_task_group() as tg:
             response = EventSourceResponse(
-                content=sse_stream_reader, data_sender_callable=sse_writer
+                content=sse_stream_reader,
+                data_sender_callable=sse_writer,
+                background=BackgroundTask(_remove_stream_writer),
             )
             logger.debug("Starting SSE response task")
             tg.start_soon(response, scope, receive, send)