Buffer per-request StreamableHTTP streams to avoid serial-router head-of-line block (#2934)

This commit is contained in:
Max 2026-06-22 16:20:45 +01:00 • committed by GitHub
commit a527142312
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 292 additions and 188 deletions

View file

@ -113,6 +113,43 @@ async def test_a_post_sse_stream_begins_with_a_priming_event_and_stamps_every_ev
)
@requirement("hosting:resume:priming")
async def test_the_priming_row_is_stored_before_any_handler_output_for_that_stream() -> None:
"""The priming cursor is the first row the event store records for a request's stream.
The POST handler stores the priming row before dispatching the request, so by construction
it precedes anything `message_router` can store for that stream id.
"""
store = SequencedEventStore()
mcp = MCPServer("resumable")
@mcp.tool()
async def burst(ctx: Context) -> str:
await ctx.info("a") # pyright: ignore[reportDeprecated]
await ctx.info("b") # pyright: ignore[reportDeprecated]
await ctx.info("c") # pyright: ignore[reportDeprecated]
return "done"
async with mounted_app(mcp, event_store=store) as (http, _):
session_id = await initialize_via_http(http)
with anyio.fail_after(5):
async with http.stream( # pragma: no branch
"POST", "/mcp", content=_tools_call(2, "burst", {}), headers=base_headers(session_id=session_id)
) as response:
await _read_events(response, 5)
# initialize wrote two rows (its own priming + response); everything after is this call.
call_rows = store._events[2:]
stream_id = call_rows[0][0]
assert [(s, None if m is None else type(m).__name__) for s, m in call_rows] == [
(stream_id, None),
(stream_id, "JSONRPCNotification"),
(stream_id, "JSONRPCNotification"),
(stream_id, "JSONRPCNotification"),
(stream_id, "JSONRPCResponse"),
]
@requirement("hosting:resume:replay")
@requirement("hosting:resume:stream-scoped")
@requirement("hosting:resume:buffered-replay")
@ -182,6 +219,46 @@ async def test_get_with_last_event_id_replays_only_that_streams_missed_events()
)
@requirement("hosting:resume:priming")
async def test_a_pre_2025_11_25_reconnect_replays_without_minting_a_priming_event() -> None:
"""A pre-2025-11-25 client reconnecting via Last-Event-ID gets the replay with no priming row.
The store-length assertion is the load-bearing proof that no priming cursor was minted.
"""
release = anyio.Event()
store = SequencedEventStore()
mcp = MCPServer("resumable")
@mcp.tool()
async def count(ctx: Context) -> str:
await ctx.info("tick 1") # pyright: ignore[reportDeprecated]
await release.wait()
await ctx.info("tick 2") # pyright: ignore[reportDeprecated]
return "counted"
async with mounted_app(mcp, event_store=store, retry_interval=0) as (http, _):
session_id = await initialize_via_http(http)
with anyio.fail_after(5):
async with http.stream(
"POST", "/mcp", content=_tools_call(1, "count", {}), headers=base_headers(session_id=session_id)
) as response:
_, first = await _read_events(response, 2)
release.set()
await store.wait_until_stored(6)
old_client_headers = base_headers(session_id=session_id) | {
"mcp-protocol-version": "2025-06-18",
"last-event-id": first.id,
}
async with http.stream("GET", "/mcp", headers=old_client_headers) as replay: # pragma: no branch
assert replay.status_code == 200
missed = await _read_events(replay, 2)
assert [(event.id, bool(event.data)) for event in missed] == snapshot([("5", True), ("6", True)])
# No priming cursor was minted on reconnect: the store still holds only the six rows
# written before the GET (init priming+response, POST priming, tick 1, tick 2, result).
assert len(store._events) == 6
@requirement("hosting:resume:bad-event-id")
async def test_an_unknown_last_event_id_yields_an_empty_replay_stream() -> None:
"""A Last-Event-ID the event store cannot map produces an empty SSE stream rather than an error.

View file

@ -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(JSONRPCResponse(jsonrpc="2.0", id="A", result={})))
await write_stream.send(SessionMessage(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

View file

@ -1634,84 +1634,6 @@ async def test_handle_sse_event_skips_empty_data() -> None:
await read_stream.aclose()
@pytest.mark.anyio
async def test_priming_event_not_sent_for_old_protocol_version() -> None:
"""_maybe_send_priming_event skips for old protocol versions (backwards compat)."""
# Create a transport with an event store
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()
@pytest.mark.anyio
async def test_priming_event_not_sent_without_event_store() -> None:
"""_maybe_send_priming_event returns early when no event_store is configured."""
# Create a transport WITHOUT an event store
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()
@pytest.mark.anyio
async def test_priming_event_includes_retry_interval() -> None:
"""_maybe_send_priming_event includes the retry field when retry_interval is set."""
# Create a transport with an event store AND retry_interval
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()
@pytest.mark.anyio
async def test_close_sse_stream_callback_not_provided_for_old_protocol_version() -> None:
"""close_sse_stream callbacks are only provided for protocol versions that support polling."""
@ -1744,30 +1666,6 @@ async def test_close_sse_stream_callback_not_provided_for_old_protocol_version()
assert session_msg_new.metadata.close_standalone_sse_stream is not None
@pytest.mark.anyio
async def test_priming_event_not_sent_for_unknown_protocol_version() -> None:
"""_maybe_send_priming_event treats unrecognized version strings conservatively.
A garbage version must not be mistaken for a future one (lexicographically
"zzz" sorts after every date-shaped revision).
"""
transport = StreamableHTTPServerTransport(
"/mcp",
event_store=SimpleEventStore(),
)
write_stream, read_stream = anyio.create_memory_object_stream[dict[str, Any]](1)
try:
await transport._maybe_send_priming_event("test-request-id", write_stream, "zzz")
# 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()
@pytest.mark.anyio
async def test_close_sse_stream_callback_not_provided_for_unknown_protocol_version() -> None:
"""close_sse_stream callbacks are withheld when the client's version is unrecognized."""