[v1.x] Buffer per-request StreamableHTTP streams; store priming event before dispatch (#2948)

src/mcp/server/streamable_http.py

@@ -14,8 +14,9 @@ from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from functools import partial
 from http import HTTPStatus
-from typing import Any
+from typing import Any, Final
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@@ -60,6 +61,11 @@ CONTENT_TYPE_SSE = "text/event-stream"
 # Special key for the standalone GET stream
 GET_STREAM_KEY = "_GET_stream"
 
+# Buffer for the per-request `_request_streams` so the serial `message_router`
+# can deposit a response and move on instead of head-of-line blocking the
+# whole session on a lazily-started `sse_writer`. See #1764.
+REQUEST_STREAM_BUFFER_SIZE: Final = 16
+
 # Session ID validation pattern (visible ASCII characters ranging from 0x21 to 0x7E)
 # Pattern ensures entire string contains only valid characters by using ^ and $ anchors
 SESSION_ID_PATTERN = re.compile(r"^[\x21-\x7E]+$")
@@ -67,6 +73,8 @@ SESSION_ID_PATTERN = re.compile(r"^[\x21-\x7E]+$")
 # Type aliases
 StreamId = str
 EventId = str
+# An SSE event-dict as accepted by sse-starlette (`event`, `data`, `id`, `retry`).
+SSEEvent = dict[str, Any]
 
 
 @dataclass
@@ -178,7 +186,7 @@ class StreamableHTTPServerTransport:
                 MemoryObjectReceiveStream[EventMessage],
             ],
         ] = {}
-        self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[dict[str, str]]] = {}
+        self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[SSEEvent]] = {}
         self._terminated = False
         # Idle timeout cancel scope; managed by the session manager.
         self.idle_scope: anyio.CancelScope | None = None
@@ -267,31 +275,48 @@ class StreamableHTTPServerTransport:
 
         return SessionMessage(message, metadata=metadata)
 
-    async def _maybe_send_priming_event(
-        self,
-        request_id: RequestId,
-        sse_stream_writer: MemoryObjectSendStream[dict[str, Any]],
-        protocol_version: str,
-    ) -> None:
-        """Send priming event for SSE resumability if event_store is configured.
+    async def _mint_priming_event(self, stream_id: StreamId, protocol_version: str) -> SSEEvent | None:
+        """Store the priming cursor for `stream_id` and return its SSE wire form.
 
-        Only sends priming events to clients with protocol version >= 2025-11-25,
-        which includes the fix for handling empty SSE data. Older clients would
-        crash trying to parse empty data as JSON.
+        Called before the request is dispatched so the priming row precedes
+        anything `message_router` can store for this stream. Returns `None`
+        when no event store is configured or the client predates 2025-11-25
+        (older clients cannot parse the empty-data event).
         """
         if not self._event_store:
-            return
-        # Priming events have empty data which older clients cannot handle.
+            return None
         if protocol_version < "2025-11-25":
-            return
-        priming_event_id = await self._event_store.store_event(
-            str(request_id),  # Convert RequestId to StreamId (str)
-            None,  # Priming event has no payload
-        )
-        priming_event: dict[str, str | int] = {"id": priming_event_id, "data": ""}
+            return None
+        priming_event_id = await self._event_store.store_event(stream_id, None)
+        priming_event: SSEEvent = {"id": priming_event_id, "data": ""}
         if self._retry_interval is not None:
             priming_event["retry"] = self._retry_interval
-        await sse_stream_writer.send(priming_event)
+        return priming_event
+
+    async def _run_sse_writer(  # pragma: no cover
+        self,
+        request_id: RequestId,
+        sse_stream_writer: MemoryObjectSendStream[SSEEvent],
+        request_stream_reader: MemoryObjectReceiveStream[EventMessage],
+        priming_event: SSEEvent | None,
+    ) -> None:
+        """Forward `_request_streams[request_id]` onto the SSE wire for one POST."""
+        try:
+            async with sse_stream_writer, request_stream_reader:
+                if priming_event is not None:
+                    await sse_stream_writer.send(priming_event)
+                async for event_message in request_stream_reader:
+                    await sse_stream_writer.send(self._create_event_data(event_message))
+                    if isinstance(event_message.message.root, JSONRPCResponse | JSONRPCError):
+                        break
+        except anyio.ClosedResourceError:
+            logger.debug("SSE stream closed by close_sse_stream()")
+        except Exception:
+            logger.exception("Error in SSE writer")
+        finally:
+            logger.debug("Closing SSE writer")
+            self._sse_stream_writers.pop(request_id, None)
+            await self._clean_up_memory_streams(request_id)
 
     def _create_error_response(
         self,
@@ -348,7 +373,7 @@ class StreamableHTTPServerTransport:
         """Extract the session ID from request headers."""
         return request.headers.get(MCP_SESSION_ID_HEADER)
 
-    def _create_event_data(self, event_message: EventMessage) -> dict[str, str]:  # pragma: no cover
+    def _create_event_data(self, event_message: EventMessage) -> SSEEvent:  # pragma: no cover
         """Create event data dictionary from an EventMessage."""
         event_data = {
             "event": "message",
@@ -530,13 +555,13 @@ class StreamableHTTPServerTransport:
                 else request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
             )
 
-            # Extract the request ID outside the try block for proper scope
-            request_id = str(message.root.id)  # pragma: no cover
-            # Register this stream for the request ID
-            self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](0)  # pragma: no cover
-            request_stream_reader = self._request_streams[request_id][1]  # pragma: no cover
+            request_id = str(message.root.id)
 
             if self.is_json_response_enabled:  # pragma: no cover
+                self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](
+                    REQUEST_STREAM_BUFFER_SIZE
+                )
+                request_stream_reader = self._request_streams[request_id][1]
                 # Process the message
                 metadata = ServerMessageMetadata(request_context=request)
                 session_message = SessionMessage(message, metadata=metadata)
@@ -580,44 +605,19 @@ class StreamableHTTPServerTransport:
                 finally:
                     await self._clean_up_memory_streams(request_id)
             else:  # pragma: no cover
-                # Create SSE stream
-                sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, str]](0)
+                # Mint the priming event before any per-request state exists:
+                # `EventStore.store_event` is user code and may raise, in which
+                # case the outer handler returns a 500 with nothing to clean up.
+                # Still strictly precedes dispatch, so storage order == wire order.
+                priming_event = await self._mint_priming_event(request_id, protocol_version)
 
-                # Store writer reference so close_sse_stream() can close it
+                sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
                 self._sse_stream_writers[request_id] = sse_stream_writer
+                self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](
+                    REQUEST_STREAM_BUFFER_SIZE
+                )
+                request_stream_reader = self._request_streams[request_id][1]
 
-                async def sse_writer():
-                    # Get the request ID from the incoming request message
-                    try:
-                        async with sse_stream_writer, request_stream_reader:
-                            # Send priming event for SSE resumability
-                            await self._maybe_send_priming_event(request_id, sse_stream_writer, protocol_version)
-
-                            # Process messages from the request-specific stream
-                            async for event_message in request_stream_reader:
-                                # Build the event data
-                                event_data = self._create_event_data(event_message)
-                                await sse_stream_writer.send(event_data)
-
-                                # If response, remove from pending streams and close
-                                if isinstance(
-                                    event_message.message.root,
-                                    JSONRPCResponse | JSONRPCError,
-                                ):
-                                    break
-                    except anyio.ClosedResourceError:
-                        # Expected when close_sse_stream() is called
-                        logger.debug("SSE stream closed by close_sse_stream()")
-                    except Exception:
-                        logger.exception("Error in SSE writer")
-                    finally:
-                        logger.debug("Closing SSE writer")
-                        self._sse_stream_writers.pop(request_id, None)
-                        await self._clean_up_memory_streams(request_id)
-
-                # Create and start EventSourceResponse
-                # SSE stream mode (original behavior)
-                # Set up headers
                 headers = {
                     "Cache-Control": "no-cache, no-transform",
                     "Connection": "keep-alive",
@@ -626,7 +626,9 @@ class StreamableHTTPServerTransport:
                 }
                 response = EventSourceResponse(
                     content=sse_stream_reader,
-                    data_sender_callable=sse_writer,
+                    data_sender_callable=partial(
+                        self._run_sse_writer, request_id, sse_stream_writer, request_stream_reader, priming_event
+                    ),
                     headers=headers,
                 )
 
@@ -644,16 +646,15 @@ class StreamableHTTPServerTransport:
                     await sse_stream_reader.aclose()
                     await self._clean_up_memory_streams(request_id)
 
-        except Exception as err:  # pragma: no cover
+        except Exception as err:
             logger.exception("Error handling POST request")
             response = self._create_error_response(
-                f"Error handling POST request: {err}",
+                "Error handling POST request",
                 HTTPStatus.INTERNAL_SERVER_ERROR,
                 INTERNAL_ERROR,
             )
             await response(scope, receive, send)
-            if writer:
-                await writer.send(Exception(err))
+            await writer.send(Exception(err))
             return
 
     async def _handle_get_request(self, request: Request, send: Send) -> None:  # pragma: no cover
@@ -706,13 +707,15 @@ class StreamableHTTPServerTransport:
             return
 
         # Create SSE stream
-        sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, str]](0)
+        sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
 
         async def standalone_sse_writer():
             try:
                 # Create a standalone message stream for server-initiated messages
 
-                self._request_streams[GET_STREAM_KEY] = anyio.create_memory_object_stream[EventMessage](0)
+                self._request_streams[GET_STREAM_KEY] = anyio.create_memory_object_stream[EventMessage](
+                    REQUEST_STREAM_BUFFER_SIZE
+                )
                 standalone_stream_reader = self._request_streams[GET_STREAM_KEY][1]
 
                 async with sse_stream_writer, standalone_stream_reader:
@@ -903,7 +906,7 @@ class StreamableHTTPServerTransport:
             replay_protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
 
             # Create SSE stream for replay
-            sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, str]](0)
+            sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
 
             async def replay_sender():
                 try:
@@ -918,22 +921,32 @@ class StreamableHTTPServerTransport:
 
                         # If stream ID not in mapping, create it
                         if stream_id and stream_id not in self._request_streams:
-                            # Register SSE writer so close_sse_stream() can close it
-                            self._sse_stream_writers[stream_id] = sse_stream_writer
+                            try:
+                                # Register SSE writer so close_sse_stream() can close it
+                                self._sse_stream_writers[stream_id] = sse_stream_writer
 
-                            # Send priming event for this new connection
-                            await self._maybe_send_priming_event(stream_id, sse_stream_writer, replay_protocol_version)
+                                # Prime the resumed connection so the client sees the stream
+                                # is re-registered. The replay→live-tail ordering window here
+                                # is pre-existing and tracked separately.
+                                priming_event = await self._mint_priming_event(stream_id, replay_protocol_version)
+                                if priming_event is not None:
+                                    await sse_stream_writer.send(priming_event)
 
-                            # Create new request streams for this connection
-                            self._request_streams[stream_id] = anyio.create_memory_object_stream[EventMessage](0)
-                            msg_reader = self._request_streams[stream_id][1]
+                                # Create new request streams for this connection
+                                self._request_streams[stream_id] = anyio.create_memory_object_stream[EventMessage](
+                                    REQUEST_STREAM_BUFFER_SIZE
+                                )
+                                msg_reader = self._request_streams[stream_id][1]
 
-                            # Forward messages to SSE
-                            async with msg_reader:
-                                async for event_message in msg_reader:
-                                    event_data = self._create_event_data(event_message)
+                                # Forward messages to SSE
+                                async with msg_reader:
+                                    async for event_message in msg_reader:
+                                        event_data = self._create_event_data(event_message)
 
-                                    await sse_stream_writer.send(event_data)
+                                        await sse_stream_writer.send(event_data)
+                            finally:
+                                self._sse_stream_writers.pop(stream_id, None)
+                                await self._clean_up_memory_streams(stream_id)
                 except anyio.ClosedResourceError:
                     # Expected when close_sse_stream() is called
                     logger.debug("Replay SSE stream closed by close_sse_stream()")

tests/server/test_streamable_http_router.py

@@ -0,0 +1,116 @@
+"""Regression coverage for the StreamableHTTP per-session response router."""
+
+import anyio
+import pytest
+from starlette.types import Message, Scope
+
+from mcp.server.streamable_http import (
+    REQUEST_STREAM_BUFFER_SIZE,
+    EventCallback,
+    EventId,
+    EventMessage,
+    EventStore,
+    StreamableHTTPServerTransport,
+    StreamId,
+)
+from mcp.shared.message import SessionMessage
+from mcp.types import JSONRPCMessage, JSONRPCResponse
+
+
+class _PrimingFailingStore(EventStore):
+    async def store_event(self, stream_id: StreamId, message: JSONRPCMessage | None) -> EventId:
+        raise RuntimeError("backend unavailable")
+
+    async def replay_events_after(self, last_event_id: EventId, send_callback: EventCallback) -> StreamId | None:
+        raise NotImplementedError
+
+
+@pytest.mark.anyio
+async def test_router_unconsumed_request_stream_does_not_block_siblings() -> None:
+    """A response whose `sse_writer` is not yet receiving must not park the router (#1764).
+
+    Drives the routing layer directly (the production race does not reproduce
+    on loopback), so this pins the router semantics, not the call sites.
+    """
+    transport = StreamableHTTPServerTransport(mcp_session_id="sid", is_json_response_enabled=False)
+    streams = transport._request_streams
+    async with transport.connect() as (_read_stream, write_stream):
+        # Model two concurrent POSTs at the point _handle_post_request has
+        # registered the per-request stream but A's sse_writer has not yet
+        # reached its first receive().
+        streams["A"] = anyio.create_memory_object_stream[EventMessage](REQUEST_STREAM_BUFFER_SIZE)
+        streams["B"] = anyio.create_memory_object_stream[EventMessage](REQUEST_STREAM_BUFFER_SIZE)
+        a_send, a_recv = streams["A"]
+        b_reader = streams["B"][1]
+        b_received = anyio.Event()
+
+        async def consume_b() -> None:
+            async with b_reader:
+                await b_reader.receive()
+                b_received.set()
+
+        async def server_writes() -> None:
+            await write_stream.send(SessionMessage(JSONRPCMessage(JSONRPCResponse(jsonrpc="2.0", id="A", result={}))))
+            await write_stream.send(SessionMessage(JSONRPCMessage(JSONRPCResponse(jsonrpc="2.0", id="B", result={}))))
+
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(consume_b)
+            tg.start_soon(server_writes)
+            with anyio.fail_after(5):
+                await b_received.wait()
+            # A's response was buffered for its (late) consumer, not dropped.
+            assert a_send.statistics().current_buffer_used == 1
+            await a_recv.aclose()
+            await a_send.aclose()
+
+
+@pytest.mark.anyio
+async def test_priming_store_failure_leaves_no_per_request_state() -> None:
+    """`EventStore.store_event` raising on the priming row must not leak per-request entries."""
+    transport = StreamableHTTPServerTransport(
+        mcp_session_id=None,
+        is_json_response_enabled=False,
+        event_store=_PrimingFailingStore(),
+    )
+
+    body = b'{"jsonrpc":"2.0","id":"req-1","method":"tools/list","params":{}}'
+    scope: Scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/",
+        "query_string": b"",
+        "headers": [
+            (b"accept", b"application/json, text/event-stream"),
+            (b"content-type", b"application/json"),
+            (b"mcp-protocol-version", b"2025-11-25"),
+        ],
+    }
+    body_sent = False
+
+    async def receive() -> Message:
+        nonlocal body_sent
+        if not body_sent:
+            body_sent = True
+            return {"type": "http.request", "body": body, "more_body": False}
+        raise NotImplementedError
+
+    sent: list[Message] = []
+
+    async def asgi_send(message: Message) -> None:
+        sent.append(message)
+
+    async with transport.connect() as (read_stream, _write_stream):
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(transport.handle_request, scope, receive, asgi_send)
+            with anyio.fail_after(5):
+                forwarded = await read_stream.receive()
+            assert isinstance(forwarded, Exception)
+        # handle_request has returned; connect()'s finally (which clears
+        # _request_streams unconditionally) has not yet run.
+        assert transport._request_streams == {}
+        assert transport._sse_stream_writers == {}
+
+    assert sent[0]["type"] == "http.response.start"
+    assert sent[0]["status"] == 500
+    body = b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body")
+    assert b"backend unavailable" not in body

tests/shared/test_streamable_http.py

@@ -1786,81 +1786,40 @@ async def test_handle_sse_event_skips_empty_data():
 
 
 @pytest.mark.anyio
-async def test_priming_event_not_sent_for_old_protocol_version():
-    """Test that _maybe_send_priming_event skips for old protocol versions (backwards compat)."""
-    # Create a transport with an event store
+async def test_priming_event_not_minted_for_old_protocol_version():
+    """`_mint_priming_event` returns None for pre-2025-11-25 clients (backwards compat)."""
     transport = StreamableHTTPServerTransport(
         "/mcp",
         event_store=SimpleEventStore(),
     )
 
-    # Create a mock stream writer
-    write_stream, read_stream = anyio.create_memory_object_stream[dict[str, Any]](1)
-
-    try:
-        # Call _maybe_send_priming_event with OLD protocol version - should NOT send
-        await transport._maybe_send_priming_event("test-request-id", write_stream, "2025-06-18")
-
-        # Nothing should have been written to the stream
-        assert write_stream.statistics().current_buffer_used == 0
-
-        # Now test with NEW protocol version - should send
-        await transport._maybe_send_priming_event("test-request-id-2", write_stream, "2025-11-25")
-
-        # Should have written a priming event
-        assert write_stream.statistics().current_buffer_used == 1
-    finally:
-        await write_stream.aclose()
-        await read_stream.aclose()
+    assert await transport._mint_priming_event("test-request-id", "2025-06-18") is None
+    event = await transport._mint_priming_event("test-request-id-2", "2025-11-25")
+    assert event is not None
+    assert event["data"] == ""
+    assert "retry" not in event
 
 
 @pytest.mark.anyio
-async def test_priming_event_not_sent_without_event_store():
-    """Test that _maybe_send_priming_event returns early when no event_store is configured."""
-    # Create a transport WITHOUT an event store
+async def test_priming_event_not_minted_without_event_store():
+    """`_mint_priming_event` returns None when no event store is configured."""
     transport = StreamableHTTPServerTransport("/mcp")
 
-    # Create a mock stream writer
-    write_stream, read_stream = anyio.create_memory_object_stream[dict[str, Any]](1)
-
-    try:
-        # Call _maybe_send_priming_event - should return early without sending
-        await transport._maybe_send_priming_event("test-request-id", write_stream, "2025-11-25")
-
-        # Nothing should have been written to the stream
-        assert write_stream.statistics().current_buffer_used == 0
-    finally:
-        await write_stream.aclose()
-        await read_stream.aclose()
+    assert await transport._mint_priming_event("test-request-id", "2025-11-25") is None
 
 
 @pytest.mark.anyio
 async def test_priming_event_includes_retry_interval():
-    """Test that _maybe_send_priming_event includes retry field when retry_interval is set."""
-    # Create a transport with an event store AND retry_interval
+    """`_mint_priming_event` carries the configured `retry` field."""
     transport = StreamableHTTPServerTransport(
         "/mcp",
         event_store=SimpleEventStore(),
         retry_interval=5000,
     )
 
-    # Create a mock stream writer
-    write_stream, read_stream = anyio.create_memory_object_stream[dict[str, Any]](1)
-
-    try:
-        # Call _maybe_send_priming_event with new protocol version
-        await transport._maybe_send_priming_event("test-request-id", write_stream, "2025-11-25")
-
-        # Should have written a priming event with retry field
-        assert write_stream.statistics().current_buffer_used == 1
-
-        # Read the event and verify it has retry field
-        event = await read_stream.receive()
-        assert "retry" in event
-        assert event["retry"] == 5000
-    finally:
-        await write_stream.aclose()
-        await read_stream.aclose()
+    event = await transport._mint_priming_event("test-request-id", "2025-11-25")
+    assert event is not None
+    assert event["retry"] == 5000
 
 
 @pytest.mark.anyio