Fix clean up memory when sse disconnect

Author: Hoa Lam

src/mcp/server/sse.py

@@ -145,7 +145,7 @@ async def sse_writer():
 
         async with anyio.create_task_group() as tg:
 
-            async def response_wrapper(scope: Scope, receive: Receive, send: Send):
+            async def response_wrapper(scope: Scope, receive: Receive, send: Send, transport: SseServerTransport):
                 """
                 The EventSourceResponse returning signals a client close / disconnect.
                 In this case we close our side of the streams to signal the client that
@@ -156,10 +156,13 @@ async def response_wrapper(scope: Scope, receive: Receive, send: Send):
                 )(scope, receive, send)
                 await read_stream_writer.aclose()
                 await write_stream_reader.aclose()
+                await read_stream.aclose()
+                await write_stream.aclose()
+                transport._read_stream_writers.pop(session_id)
                 logging.debug(f"Client session disconnected {session_id}")
 
             logger.debug("Starting SSE response task")
-            tg.start_soon(response_wrapper, scope, receive, send)
+            tg.start_soon(response_wrapper, scope, receive, send, self)
 
             logger.debug("Yielding read and write streams")
             yield (read_stream, write_stream)

tests/server/test_sse_disconnect.py

@@ -46,9 +46,8 @@ async def send(message: Message) -> None:
         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
+    # Check that the session_id should be clean up
+    assert session_id not in transport._read_stream_writers
 
     # After context exits, session should be cleaned up
     assert len(transport._read_stream_writers) == 0