[v1.x] Buffer per-request StreamableHTTP streams; store priming event before dispatch (#2948)
This commit is contained in:
Max
GitHub
parent
32d32908fe
commit
6df3d73426
2 changed files with 229 additions and 141 deletions
116
tests/server/test_streamable_http_router.py
Normal file
116
tests/server/test_streamable_http_router.py
Normal 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(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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue