Fix: memory leak and exit problems

examples/servers/simple-prompt/mcp_simple_prompt/server.py

@@ -90,17 +90,23 @@ async def get_prompt(
     if transport == "sse":
         from mcp.server.sse import SseServerTransport
         from starlette.applications import Starlette
+        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()
-                )
+            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,

examples/servers/simple-resource/mcp_simple_resource/server.py

@@ -46,17 +46,23 @@ async def read_resource(uri: FileUrl) -> str | bytes:
     if transport == "sse":
         from mcp.server.sse import SseServerTransport
         from starlette.applications import Starlette
+        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()
-                )
+            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,

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,

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/fastmcp/server.py

@@ -20,6 +20,7 @@
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from starlette.applications import Starlette
 from starlette.requests import Request
+from starlette.responses import Response
 from starlette.routing import Mount, Route
 
 from mcp.server.fastmcp.exceptions import ResourceError
@@ -482,17 +483,20 @@ def sse_app(self) -> Starlette:
         """Return an instance of the SSE server app."""
         sse = SseServerTransport(self.settings.message_path)
 
-        async def handle_sse(request: Request) -> None:
-            async with sse.connect_sse(
-                request.scope,
-                request.receive,
-                request._send,  # type: ignore[reportPrivateUsage]
-            ) as streams:
-                await self._mcp_server.run(
-                    streams[0],
-                    streams[1],
-                    self._mcp_server.create_initialization_options(),
-                )
+        async def handle_sse(request: Request) -> Response:
+            with anyio.CancelScope() as cancel_scope:
+                async with sse.connect_sse(
+                    request.scope,
+                    request.receive,
+                    request._send,  # type: ignore[reportPrivateUsage]
+                    callback=lambda: cancel_scope.cancel(),
+                ) as streams:
+                    await self._mcp_server.run(
+                        streams[0],
+                        streams[1],
+                        self._mcp_server.create_initialization_options(),
+                    )
+            return Response(status_code=200)
 
         return Starlette(
             debug=self.settings.debug,

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
@@ -43,7 +44,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
 
@@ -79,7 +80,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 +127,19 @@ async def sse_writer():
                         }
                     )
 
+        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]
+            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,
+                client_close_handler_callable=client_close_handler,  # type: ignore[arg-type]
             )
             logger.debug("Starting SSE response task")
             tg.start_soon(response, scope, receive, send)

tests/shared/test_sse.py

@@ -10,6 +10,7 @@
 from pydantic import AnyUrl
 from starlette.applications import Starlette
 from starlette.requests import Request
+from starlette.responses import Response
 from starlette.routing import Mount, Route
 
 from mcp.client.session import ClientSession
@@ -83,13 +84,18 @@ def make_server_app() -> Starlette:
     sse = SseServerTransport("/messages/")
     server = ServerTest()
 
-    async def handle_sse(request: Request) -> None:
-        async with sse.connect_sse(
-            request.scope, request.receive, request._send
-        ) as streams:
-            await server.run(
-                streams[0], streams[1], server.create_initialization_options()
-            )
+    async def handle_sse(request: Request) -> Response:
+        with anyio.CancelScope() as cancel_scope:
+            async with sse.connect_sse(
+                request.scope,
+                request.receive,
+                request._send,
+                lambda: cancel_scope.cancel(),
+            ) as streams:
+                await server.run(
+                    streams[0], streams[1], server.create_initialization_options()
+                )
+        return Response(status_code=200)
 
     app = Starlette(
         routes=[