feat: graceful SSE drain on session manager shutdown

Terminate all active transports before cancelling the task group during
StreamableHTTPSessionManager shutdown.  This closes their in-memory
streams, allowing EventSourceResponse to send a final `more_body=False`
chunk — a clean HTTP close instead of a connection reset.

Without this, reverse proxies like nginx see "upstream prematurely
closed connection" and return 502 to clients during rolling deploys.

Changes:
- Track in-flight stateless transports in `_stateless_transports` set
- In `run()` finally block, call `terminate()` on all stateful and
  stateless transports before `tg.cancel_scope.cancel()`
- Add E2E tests for both stateless and stateful modes that verify
  the SSE stream closes cleanly when the manager shuts down while
  a tool call is in-flight

Author: wiggzz

src/mcp/server/streamable_http_manager.py

@@ -90,6 +90,9 @@ def __init__(
         self._session_creation_lock = anyio.Lock()
         self._server_instances: dict[str, StreamableHTTPServerTransport] = {}
 
+        # Track in-flight stateless transports for graceful shutdown
+        self._stateless_transports: set[StreamableHTTPServerTransport] = set()
+
         # The task group will be set during lifespan
         self._task_group = None
         # Thread-safe tracking of run() calls
@@ -130,11 +133,28 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]:
                 yield  # Let the application run
             finally:
                 logger.info("StreamableHTTP session manager shutting down")
+
+                # Terminate all active transports before cancelling the task
+                # group.  This closes their in-memory streams, which lets
+                # EventSourceResponse send a final ``more_body=False`` chunk
+                # — a clean HTTP close instead of a connection reset.
+                for transport in list(self._server_instances.values()):
+                    try:
+                        await transport.terminate()
+                    except Exception:  # pragma: no cover
+                        logger.debug("Error terminating transport during shutdown", exc_info=True)
+                for transport in list(self._stateless_transports):
+                    try:
+                        await transport.terminate()
+                    except Exception:  # pragma: no cover
+                        logger.debug("Error terminating stateless transport during shutdown", exc_info=True)
+
                 # Cancel task group to stop all spawned tasks
                 tg.cancel_scope.cancel()
                 self._task_group = None
                 # Clear any remaining server instances
                 self._server_instances.clear()
+                self._stateless_transports.clear()
 
     async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Process ASGI request with proper session handling and transport setup.
@@ -161,6 +181,9 @@ async def _handle_stateless_request(self, scope: Scope, receive: Receive, send:
             security_settings=self.security_settings,
         )
 
+        # Track for graceful shutdown
+        self._stateless_transports.add(http_transport)
+
         # Start server in a new task
         async def run_stateless_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED):
             async with http_transport.connect() as streams:
@@ -181,8 +204,11 @@ async def run_stateless_server(*, task_status: TaskStatus[None] = anyio.TASK_STA
         # Start the server task
         await self._task_group.start(run_stateless_server)
 
-        # Handle the HTTP request and return the response
-        await http_transport.handle_request(scope, receive, send)
+        try:
+            # Handle the HTTP request and return the response
+            await http_transport.handle_request(scope, receive, send)
+        finally:
+            self._stateless_transports.discard(http_transport)
 
         # Terminate the transport after the request is handled
         await http_transport.terminate()

tests/server/test_streamable_http_manager.py

@@ -10,11 +10,11 @@
 import pytest
 from starlette.types import Message
 
-from mcp import Client
+from mcp import Client, types
 from mcp.client.streamable_http import streamable_http_client
 from mcp.server import Server, ServerRequestContext, streamable_http_manager
 from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, StreamableHTTPServerTransport
-from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
 from mcp.types import INVALID_REQUEST, ListToolsResult, PaginatedRequestParams
 
 
@@ -413,3 +413,204 @@ def test_session_idle_timeout_rejects_non_positive():
 def test_session_idle_timeout_rejects_stateless():
     with pytest.raises(RuntimeError, match="not supported in stateless"):
         StreamableHTTPSessionManager(app=Server("test"), session_idle_timeout=30, stateless=True)
+
+
+MCP_HEADERS = {
+    "Accept": "application/json, text/event-stream",
+    "Content-Type": "application/json",
+}
+
+_INITIALIZE_REQUEST = {
+    "jsonrpc": "2.0",
+    "id": 1,
+    "method": "initialize",
+    "params": {
+        "protocolVersion": "2025-03-26",
+        "capabilities": {},
+        "clientInfo": {"name": "test", "version": "0.1"},
+    },
+}
+
+_INITIALIZED_NOTIFICATION = {
+    "jsonrpc": "2.0",
+    "method": "notifications/initialized",
+}
+
+_TOOL_CALL_REQUEST = {
+    "jsonrpc": "2.0",
+    "id": 2,
+    "method": "tools/call",
+    "params": {"name": "slow_tool", "arguments": {"message": "hello"}},
+}
+
+
+def _make_slow_tool_server() -> tuple[Server, anyio.Event]:
+    """Create an MCP server with a tool that blocks forever, returning
+    the server and an event that fires when the tool starts executing."""
+    tool_started = anyio.Event()
+
+    async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
+        tool_started.set()
+        await anyio.sleep_forever()
+        return types.CallToolResult(  # pragma: no cover
+            content=[types.TextContent(type="text", text="never reached")]
+        )
+
+    async def handle_list_tools(
+        ctx: ServerRequestContext, params: PaginatedRequestParams | None
+    ) -> ListToolsResult:  # pragma: no cover
+        return ListToolsResult(
+            tools=[
+                types.Tool(
+                    name="slow_tool",
+                    description="A tool that blocks forever",
+                    input_schema={"type": "object", "properties": {"message": {"type": "string"}}},
+                )
+            ]
+        )
+
+    app = Server("test-graceful-shutdown", on_call_tool=handle_call_tool, on_list_tools=handle_list_tools)
+    return app, tool_started
+
+
+@pytest.mark.anyio
+async def test_graceful_shutdown_terminates_active_stateless_transports():
+    """Verify that shutting down the session manager terminates in-flight
+    stateless transports so SSE streams close cleanly (``more_body=False``)
+    instead of being abruptly cancelled.
+
+    This prevents "upstream prematurely closed connection" errors at reverse
+    proxies like nginx.
+    """
+    app, tool_started = _make_slow_tool_server()
+    manager = StreamableHTTPSessionManager(app=app, stateless=True)
+
+    mcp_app = StreamableHTTPASGIApp(manager)
+
+    manager_ready = anyio.Event()
+    shutdown_event = anyio.Event()
+    stream_outcome: str | None = None
+
+    async with anyio.create_task_group() as tg:
+
+        async def run_lifespan_and_shutdown():
+            async with manager.run():
+                manager_ready.set()
+                with anyio.fail_after(5):
+                    await tool_started.wait()
+            shutdown_event.set()
+
+        async def make_requests():
+            nonlocal stream_outcome
+            with anyio.fail_after(5):
+                await manager_ready.wait()
+            async with (
+                httpx.ASGITransport(mcp_app) as transport,
+                httpx.AsyncClient(transport=transport, base_url="http://testserver") as client,
+            ):
+                # Initialize
+                resp = await client.post("/mcp/", json=_INITIALIZE_REQUEST, headers=MCP_HEADERS)
+                resp.raise_for_status()
+
+                # Send initialized notification
+                resp = await client.post("/mcp/", json=_INITIALIZED_NOTIFICATION, headers=MCP_HEADERS)
+                assert resp.status_code == 202
+
+                # Send slow tool call — this returns an SSE stream
+                try:
+                    async with client.stream(
+                        "POST",
+                        "/mcp/",
+                        json=_TOOL_CALL_REQUEST,
+                        headers=MCP_HEADERS,
+                        timeout=httpx.Timeout(10, connect=5),
+                    ) as stream:
+                        stream.raise_for_status()
+                        async for _chunk in stream.aiter_bytes():
+                            pass  # pragma: no cover
+                    stream_outcome = "clean"
+                except httpx.RemoteProtocolError:  # pragma: no cover
+                    stream_outcome = "reset"
+
+        tg.start_soon(run_lifespan_and_shutdown)
+        tg.start_soon(make_requests)
+
+        with anyio.fail_after(10):
+            await shutdown_event.wait()
+
+        tg.cancel_scope.cancel()
+
+    assert stream_outcome == "clean", f"Expected clean HTTP close, got {stream_outcome}"
+
+
+@pytest.mark.anyio
+async def test_graceful_shutdown_terminates_active_stateful_transports():
+    """Verify that shutting down the session manager terminates in-flight
+    stateful transports so SSE streams close cleanly."""
+    app, tool_started = _make_slow_tool_server()
+    manager = StreamableHTTPSessionManager(app=app, stateless=False)
+
+    mcp_app = StreamableHTTPASGIApp(manager)
+
+    manager_ready = anyio.Event()
+    shutdown_event = anyio.Event()
+    stream_outcome: str | None = None
+
+    async with anyio.create_task_group() as tg:
+
+        async def run_lifespan_and_shutdown():
+            async with manager.run():
+                manager_ready.set()
+                with anyio.fail_after(5):
+                    await tool_started.wait()
+            shutdown_event.set()
+
+        async def make_requests():
+            nonlocal stream_outcome
+            with anyio.fail_after(5):
+                await manager_ready.wait()
+            async with (
+                httpx.ASGITransport(mcp_app) as transport,
+                httpx.AsyncClient(transport=transport, base_url="http://testserver") as client,
+            ):
+                # Initialize (creates a session)
+                resp = await client.post("/mcp/", json=_INITIALIZE_REQUEST, headers=MCP_HEADERS)
+                resp.raise_for_status()
+                session_id = resp.headers.get(MCP_SESSION_ID_HEADER)
+                assert session_id is not None
+
+                session_headers = {
+                    **MCP_HEADERS,
+                    MCP_SESSION_ID_HEADER: session_id,
+                    "mcp-protocol-version": "2025-03-26",
+                }
+
+                # Send initialized notification
+                resp = await client.post("/mcp/", json=_INITIALIZED_NOTIFICATION, headers=session_headers)
+                assert resp.status_code == 202
+
+                # Send slow tool call
+                try:
+                    async with client.stream(
+                        "POST",
+                        "/mcp/",
+                        json=_TOOL_CALL_REQUEST,
+                        headers=session_headers,
+                        timeout=httpx.Timeout(10, connect=5),
+                    ) as stream:
+                        stream.raise_for_status()
+                        async for _chunk in stream.aiter_bytes():
+                            pass  # pragma: no cover
+                    stream_outcome = "clean"
+                except httpx.RemoteProtocolError:  # pragma: no cover
+                    stream_outcome = "reset"
+
+        tg.start_soon(run_lifespan_and_shutdown)
+        tg.start_soon(make_requests)
+
+        with anyio.fail_after(10):
+            await shutdown_event.wait()
+
+        tg.cancel_scope.cancel()
+
+    assert stream_outcome == "clean", f"Expected clean HTTP close, got {stream_outcome}"