fix(server): terminate streamable HTTP sessions on manager shutdown

Why: On CTRL+C / lifespan exit with live SSE clients, StreamableHTTPSessionManager
cancelled the task group without terminating transports, and terminate() left
_sse_stream_writers open — uvicorn then logs "ASGI callable returned without
completing response" (#2150).

- Manager run() finally: await terminate() on active transports before cancel
- terminate(): close all SSE writers first (incl. GET + provisional replay keys)
- Register GET standalone writer immediately; provisional key during replay
- Regression tests for writer close + manager shutdown path

Rebased onto current main after #3125 was closed unmerged during v2 backlog
cleanup (maintainers pointed at #2253 on v1.x; this targets main with tests).

Signed-off-by: Atlas/luccc-grok-4.5 <atlas@agent-studio.local>

Author: OpenSourceMaintenance-Luccc-grok-4.5

src/mcp/server/streamable_http.py

@@ -734,6 +734,8 @@ async def _handle_get_request(self, request: Request, send: Send) -> None:
 
         # Create SSE stream
         sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
+        # Register immediately so session terminate() can close this ASGI response.
+        self._sse_stream_writers[GET_STREAM_KEY] = sse_stream_writer
 
         async def standalone_sse_writer():
             try:
@@ -757,11 +759,13 @@ async def standalone_sse_writer():
                         await sse_stream_writer.send(event_data)
             except anyio.ClosedResourceError:
                 # Session teardown can close the stream while the writer is between dequeues.
+                # Also expected when terminate()/close_sse_stream() closes the writer.
                 pass
             except Exception:
                 logger.exception("Error in standalone SSE writer")  # pragma: no cover
             finally:
                 logger.debug("Closing standalone SSE writer")
+                self._sse_stream_writers.pop(GET_STREAM_KEY, None)
                 await self._clean_up_memory_streams(GET_STREAM_KEY)
 
         # Create and start EventSourceResponse
@@ -778,6 +782,7 @@ async def standalone_sse_writer():
             logger.exception("Error in standalone SSE response")
             await self._clean_up_memory_streams(GET_STREAM_KEY)
         finally:
+            self._sse_stream_writers.pop(GET_STREAM_KEY, None)
             await sse_stream_writer.aclose()
             await sse_stream_reader.aclose()
 
@@ -808,11 +813,19 @@ async def terminate(self) -> None:
         """Terminate the current session, closing all streams.
 
         Once terminated, all requests with this session ID will receive 404 Not Found.
+
+        Active SSE writers are closed first so EventSourceResponse / ASGI callables can
+        complete instead of hanging until the task group is cancelled (see #2150).
         """
 
         self._terminated = True
         logger.info(f"Terminating session: {self.mcp_session_id}")
 
+        # Close SSE stream writers first so long-lived GET/POST SSE responses finish.
+        # Copy keys: close_sse_stream mutates the dict (includes GET_STREAM_KEY).
+        for request_id in list(self._sse_stream_writers.keys()):
+            self.close_sse_stream(request_id)
+
         # We need a copy of the keys to avoid modification during iteration
         request_stream_keys = list(self._request_streams.keys())
 
@@ -910,8 +923,13 @@ async def _replay_events(self, last_event_id: str, request: Request, send: Send)
 
             # Create SSE stream for replay
             sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
+            # Provisional key so terminate() can close the response during replay itself
+            # (before stream_id is known from event_store.replay_events_after).
+            replay_writer_key: RequestId = f"_replay:{last_event_id}:{id(sse_stream_writer)}"
+            self._sse_stream_writers[replay_writer_key] = sse_stream_writer
 
             async def replay_sender():
+                registered_stream_id: RequestId | None = None
                 try:
                     async with sse_stream_writer:
                         # Define an async callback for sending events
@@ -925,8 +943,10 @@ async def send_event(event_message: EventMessage) -> None:
                         # If stream ID not in mapping, create it
                         if stream_id and stream_id not in self._request_streams:  # pragma: no branch
                             try:
-                                # Register SSE writer so close_sse_stream() can close it
+                                # Re-key from provisional → stream_id for close_sse_stream(stream_id)
+                                self._sse_stream_writers.pop(replay_writer_key, None)
                                 self._sse_stream_writers[stream_id] = sse_stream_writer
+                                registered_stream_id = stream_id
 
                                 # Prime the resumed connection so the client sees the stream
                                 # is re-registered. The replay→live-tail ordering window here
@@ -948,13 +968,20 @@ async def send_event(event_message: EventMessage) -> None:
 
                                         await sse_stream_writer.send(event_data)
                             finally:
-                                self._sse_stream_writers.pop(stream_id, None)
+                                # registered_stream_id is set immediately on try entry; keep the
+                                # guard for defensive cleanup if re-key is later reordered.
+                                if registered_stream_id is not None:  # pragma: no branch
+                                    self._sse_stream_writers.pop(registered_stream_id, None)
                                 await self._clean_up_memory_streams(stream_id)
                 except anyio.ClosedResourceError:  # pragma: lax no cover
                     # Expected when close_sse_stream() is called
                     logger.debug("Replay SSE stream closed by close_sse_stream()")
                 except Exception:  # pragma: lax no cover
                     logger.exception("Error in replay sender")
+                finally:
+                    self._sse_stream_writers.pop(replay_writer_key, None)
+                    if registered_stream_id is not None:
+                        self._sse_stream_writers.pop(registered_stream_id, None)
 
             # Create and start EventSourceResponse
             response = EventSourceResponse(
@@ -968,6 +995,7 @@ async def send_event(event_message: EventMessage) -> None:
             except Exception:  # pragma: lax no cover
                 logger.exception("Error in replay response")
             finally:
+                self._sse_stream_writers.pop(replay_writer_key, None)
                 await sse_stream_writer.aclose()
                 await sse_stream_reader.aclose()
 

src/mcp/server/streamable_http_manager.py

@@ -153,6 +153,18 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]:
                 yield  # Let the application run
             finally:
                 logger.info("StreamableHTTP session manager shutting down")
+                # Terminate active transports before cancelling the task group so
+                # in-flight SSE responses can complete cleanly (issue #2150).
+                active_transports = list(self._server_instances.values())
+                for transport in active_transports:
+                    if not transport.is_terminated:  # pragma: no branch
+                        try:
+                            await transport.terminate()
+                        except Exception:  # pragma: no cover
+                            logger.exception(
+                                "Error terminating streamable HTTP session %s during shutdown",
+                                transport.mcp_session_id,
+                            )
                 # Cancel task group to stop all spawned tasks
                 tg.cancel_scope.cancel()
                 self._task_group = None

tests/server/test_streamable_http_manager.py

@@ -746,3 +746,90 @@ async def test_anonymous_session_accepts_anonymous_requests(
     session_id = await _open_session(manager, None)
 
     assert await _request_session(manager, session_id, None) != 404
+
+
+@pytest.mark.anyio
+async def test_terminate_closes_active_sse_stream_writers():
+    """Regression for #2150: terminate must close SSE writers so ASGI can finish.
+
+    Without this, manager shutdown cancels the task group while EventSourceResponse
+    is still open and uvicorn logs "ASGI callable returned without completing response".
+    """
+    transport = StreamableHTTPServerTransport(mcp_session_id="test-session-2150")
+    send_stream, receive_stream = anyio.create_memory_object_stream[object](1)
+    transport._sse_stream_writers["req-1"] = send_stream  # type: ignore[assignment]
+    # Orphaned request stream with a key NOT also in _sse_stream_writers: close_sse_stream
+    # only pops matching request streams, so terminate's remaining-keys loop must run.
+    req_send, req_recv = anyio.create_memory_object_stream[object](1)
+    transport._request_streams["req-orphan"] = (req_send, req_recv)  # type: ignore[assignment]
+
+    await transport.terminate()
+
+    assert transport.is_terminated
+    assert "req-1" not in transport._sse_stream_writers
+    assert not transport._request_streams
+    with pytest.raises(anyio.ClosedResourceError):
+        await send_stream.send(object())  # type: ignore[arg-type]
+    await receive_stream.aclose()
+
+
+@pytest.mark.anyio
+async def test_manager_shutdown_terminates_active_sessions():
+    """Regression for #2150: run() finally should terminate tracked transports."""
+    app = Server("test-shutdown-terminate")
+    manager = StreamableHTTPSessionManager(app=app)
+    transport = StreamableHTTPServerTransport(mcp_session_id="shutdown-session")
+    # Inject a live session as if a client still held an SSE connection.
+    manager._server_instances[transport.mcp_session_id] = transport  # type: ignore[index]
+    original_terminate = transport.terminate
+    terminate_calls = 0
+
+    async def counting_terminate() -> None:
+        nonlocal terminate_calls
+        terminate_calls += 1
+        await original_terminate()
+
+    transport.terminate = counting_terminate  # type: ignore[method-assign]
+
+    async with manager.run():
+        assert transport.mcp_session_id in manager._server_instances
+        # Exit context -> shutdown path should terminate then clear.
+
+    assert terminate_calls == 1
+    assert transport.is_terminated
+    assert transport.mcp_session_id not in manager._server_instances
+    assert not manager._server_instances
+
+
+@pytest.mark.anyio
+async def test_terminate_closes_standalone_get_sse_writer_when_registered():
+    """GET standalone SSE writers must be registered so terminate can close them (#2150)."""
+    from mcp.server.streamable_http import GET_STREAM_KEY
+
+    transport = StreamableHTTPServerTransport(mcp_session_id="get-session-2150")
+    send_stream, receive_stream = anyio.create_memory_object_stream[object](1)
+    transport._sse_stream_writers[GET_STREAM_KEY] = send_stream  # type: ignore[assignment]
+
+    await transport.terminate()
+
+    assert transport.is_terminated
+    assert GET_STREAM_KEY not in transport._sse_stream_writers
+    with pytest.raises(anyio.ClosedResourceError):
+        await send_stream.send(object())  # type: ignore[arg-type]
+    await receive_stream.aclose()
+
+
+@pytest.mark.anyio
+async def test_terminate_closes_provisional_replay_sse_writer():
+    """Replay writers registered under provisional keys must close on terminate."""
+    transport = StreamableHTTPServerTransport(mcp_session_id="replay-session-2150")
+    send_stream, receive_stream = anyio.create_memory_object_stream[object](1)
+    key = "_replay:evt-1:123"
+    transport._sse_stream_writers[key] = send_stream  # type: ignore[assignment]
+
+    await transport.terminate()
+
+    assert key not in transport._sse_stream_writers
+    with pytest.raises(anyio.ClosedResourceError):
+        await send_stream.send(object())  # type: ignore[arg-type]
+    await receive_stream.aclose()