Fix handle sse disconnect event to free read_stream_writers[session_id]

src/mcp/server/sse.py

@@ -43,7 +43,7 @@ async def handle_sse(request):
 from sse_starlette import EventSourceResponse
 from starlette.requests import Request
 from starlette.responses import Response
-from starlette.types import Receive, Scope, Send
+from starlette.types import Message, Receive, Scope, Send
 
 import mcp.types as types
 
@@ -120,9 +120,19 @@ async def sse_writer():
                         }
                     )
 
+        async def handle_see_disconnect(message: Message) -> None:
+            logger.debug(f"Disconnect sse {session_id}")
+            del self._read_stream_writers[session_id]
+            await read_stream.aclose()
+            await read_stream_writer.aclose()
+            await write_stream.aclose()
+            await write_stream_reader.aclose()
+
         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,
+                client_close_handler_callable=handle_see_disconnect  # type: ignore
             )
             logger.debug("Starting SSE response task")
             tg.start_soon(response, scope, receive, send)

tests/server/test_sse_disconnect.py

@@ -0,0 +1,54 @@
+import asyncio
+from uuid import UUID
+
+import pytest
+from starlette.types import Message, Scope
+
+from mcp.server.sse import SseServerTransport
+
+
+@pytest.mark.anyio
+async def test_sse_disconnect_handle():
+    transport = SseServerTransport(endpoint="/sse")
+    # Create a minimal ASGI scope for an HTTP GET request
+    scope: Scope = {
+        "type": "http",
+        "method": "GET",
+        "path": "/sse",
+        "headers": [],
+    }
+    send_disconnect = False
+
+    # Dummy receive and send functions
+    async def receive() -> dict:
+        nonlocal send_disconnect
+        if not send_disconnect:
+            send_disconnect = True
+            return {"type": "http.request"}
+        else:
+            return {"type": "http.disconnect"}
+
+    async def send(message: Message) -> None:
+        await asyncio.sleep(0)
+
+    # Run the connect_sse context manager
+    async with transport.connect_sse(scope, receive, send) as (
+        read_stream,
+        write_stream,
+    ):
+        # Assert that streams are provided
+        assert read_stream is not None
+        assert write_stream is not None
+
+        # There should be exactly one session
+        assert len(transport._read_stream_writers) == 1
+        # Check that the session key is a UUID
+        session_id = next(iter(transport._read_stream_writers.keys()))
+        assert isinstance(session_id, UUID)
+
+        # Check that the writer is still open
+        writer = transport._read_stream_writers[session_id]
+        assert writer is not None
+
+    # After context exits, session should be cleaned up
+    assert len(transport._read_stream_writers) == 0