use client_close_handler_callable for connect_sse

Author: PWZER

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "httpx-sse>=0.4",
     "pydantic>=2.7.2,<3.0.0",
     "starlette>=0.27",
-    "sse-starlette>=1.6.1",
+    "sse-starlette>=2.3.0",
     "pydantic-settings>=2.5.2",
     "uvicorn>=0.23.1; sys_platform != 'emscripten'",
 ]

src/mcp/server/sse.py

@@ -42,10 +42,9 @@ 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
+from starlette.types import Receive, Scope, Send, Message
 
 import mcp.types as types
 
@@ -128,7 +127,7 @@ async def sse_writer():
                         }
                     )
 
-        async def _remove_stream_writer() -> None:
+        async def client_close_handler(message: Message) -> None:
             await read_stream_writer.aclose()
             await write_stream_reader.aclose()
             del self._read_stream_writers[session_id]
@@ -140,7 +139,7 @@ async def _remove_stream_writer() -> None:
             response = EventSourceResponse(
                 content=sse_stream_reader,
                 data_sender_callable=sse_writer,
-                background=BackgroundTask(_remove_stream_writer),
+                client_close_handler_callable=client_close_handler,
             )
             logger.debug("Starting SSE response task")
             tg.start_soon(response, scope, receive, send)