[v1.x] Buffer per-request StreamableHTTP streams; store priming event before dispatch (#2948)
This commit is contained in:
parent
32d32908fe
commit
6df3d73426
3 changed files with 229 additions and 141 deletions
|
|
@ -14,8 +14,9 @@ from abc import ABC, abstractmethod
|
||||||
from collections.abc import AsyncGenerator, Awaitable, Callable
|
from collections.abc import AsyncGenerator, Awaitable, Callable
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
from functools import partial
|
||||||
from http import HTTPStatus
|
from http import HTTPStatus
|
||||||
from typing import Any
|
from typing import Any, Final
|
||||||
|
|
||||||
import anyio
|
import anyio
|
||||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||||
|
|
@ -60,6 +61,11 @@ CONTENT_TYPE_SSE = "text/event-stream"
|
||||||
# Special key for the standalone GET stream
|
# Special key for the standalone GET stream
|
||||||
GET_STREAM_KEY = "_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)
|
# Session ID validation pattern (visible ASCII characters ranging from 0x21 to 0x7E)
|
||||||
# Pattern ensures entire string contains only valid characters by using ^ and $ anchors
|
# Pattern ensures entire string contains only valid characters by using ^ and $ anchors
|
||||||
SESSION_ID_PATTERN = re.compile(r"^[\x21-\x7E]+$")
|
SESSION_ID_PATTERN = re.compile(r"^[\x21-\x7E]+$")
|
||||||
|
|
@ -67,6 +73,8 @@ SESSION_ID_PATTERN = re.compile(r"^[\x21-\x7E]+$")
|
||||||
# Type aliases
|
# Type aliases
|
||||||
StreamId = str
|
StreamId = str
|
||||||
EventId = str
|
EventId = str
|
||||||
|
# An SSE event-dict as accepted by sse-starlette (`event`, `data`, `id`, `retry`).
|
||||||
|
SSEEvent = dict[str, Any]
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
@ -178,7 +186,7 @@ class StreamableHTTPServerTransport:
|
||||||
MemoryObjectReceiveStream[EventMessage],
|
MemoryObjectReceiveStream[EventMessage],
|
||||||
],
|
],
|
||||||
] = {}
|
] = {}
|
||||||
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[dict[str, str]]] = {}
|
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[SSEEvent]] = {}
|
||||||
self._terminated = False
|
self._terminated = False
|
||||||
# Idle timeout cancel scope; managed by the session manager.
|
# Idle timeout cancel scope; managed by the session manager.
|
||||||
self.idle_scope: anyio.CancelScope | None = None
|
self.idle_scope: anyio.CancelScope | None = None
|
||||||
|
|
@ -267,31 +275,48 @@ class StreamableHTTPServerTransport:
|
||||||
|
|
||||||
return SessionMessage(message, metadata=metadata)
|
return SessionMessage(message, metadata=metadata)
|
||||||
|
|
||||||
async def _maybe_send_priming_event(
|
async def _mint_priming_event(self, stream_id: StreamId, protocol_version: str) -> SSEEvent | None:
|
||||||
self,
|
"""Store the priming cursor for `stream_id` and return its SSE wire form.
|
||||||
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.
|
|
||||||
|
|
||||||
Only sends priming events to clients with protocol version >= 2025-11-25,
|
Called before the request is dispatched so the priming row precedes
|
||||||
which includes the fix for handling empty SSE data. Older clients would
|
anything `message_router` can store for this stream. Returns `None`
|
||||||
crash trying to parse empty data as JSON.
|
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:
|
if not self._event_store:
|
||||||
return
|
return None
|
||||||
# Priming events have empty data which older clients cannot handle.
|
|
||||||
if protocol_version < "2025-11-25":
|
if protocol_version < "2025-11-25":
|
||||||
return
|
return None
|
||||||
priming_event_id = await self._event_store.store_event(
|
priming_event_id = await self._event_store.store_event(stream_id, None)
|
||||||
str(request_id), # Convert RequestId to StreamId (str)
|
priming_event: SSEEvent = {"id": priming_event_id, "data": ""}
|
||||||
None, # Priming event has no payload
|
|
||||||
)
|
|
||||||
priming_event: dict[str, str | int] = {"id": priming_event_id, "data": ""}
|
|
||||||
if self._retry_interval is not None:
|
if self._retry_interval is not None:
|
||||||
priming_event["retry"] = self._retry_interval
|
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(
|
def _create_error_response(
|
||||||
self,
|
self,
|
||||||
|
|
@ -348,7 +373,7 @@ class StreamableHTTPServerTransport:
|
||||||
"""Extract the session ID from request headers."""
|
"""Extract the session ID from request headers."""
|
||||||
return request.headers.get(MCP_SESSION_ID_HEADER)
|
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."""
|
"""Create event data dictionary from an EventMessage."""
|
||||||
event_data = {
|
event_data = {
|
||||||
"event": "message",
|
"event": "message",
|
||||||
|
|
@ -530,13 +555,13 @@ class StreamableHTTPServerTransport:
|
||||||
else request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
|
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)
|
||||||
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
|
|
||||||
|
|
||||||
if self.is_json_response_enabled: # pragma: no cover
|
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
|
# Process the message
|
||||||
metadata = ServerMessageMetadata(request_context=request)
|
metadata = ServerMessageMetadata(request_context=request)
|
||||||
session_message = SessionMessage(message, metadata=metadata)
|
session_message = SessionMessage(message, metadata=metadata)
|
||||||
|
|
@ -580,44 +605,19 @@ class StreamableHTTPServerTransport:
|
||||||
finally:
|
finally:
|
||||||
await self._clean_up_memory_streams(request_id)
|
await self._clean_up_memory_streams(request_id)
|
||||||
else: # pragma: no cover
|
else: # pragma: no cover
|
||||||
# Create SSE stream
|
# Mint the priming event before any per-request state exists:
|
||||||
sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, str]](0)
|
# `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._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 = {
|
headers = {
|
||||||
"Cache-Control": "no-cache, no-transform",
|
"Cache-Control": "no-cache, no-transform",
|
||||||
"Connection": "keep-alive",
|
"Connection": "keep-alive",
|
||||||
|
|
@ -626,7 +626,9 @@ class StreamableHTTPServerTransport:
|
||||||
}
|
}
|
||||||
response = EventSourceResponse(
|
response = EventSourceResponse(
|
||||||
content=sse_stream_reader,
|
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,
|
headers=headers,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -644,16 +646,15 @@ class StreamableHTTPServerTransport:
|
||||||
await sse_stream_reader.aclose()
|
await sse_stream_reader.aclose()
|
||||||
await self._clean_up_memory_streams(request_id)
|
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")
|
logger.exception("Error handling POST request")
|
||||||
response = self._create_error_response(
|
response = self._create_error_response(
|
||||||
f"Error handling POST request: {err}",
|
"Error handling POST request",
|
||||||
HTTPStatus.INTERNAL_SERVER_ERROR,
|
HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||||
INTERNAL_ERROR,
|
INTERNAL_ERROR,
|
||||||
)
|
)
|
||||||
await response(scope, receive, send)
|
await response(scope, receive, send)
|
||||||
if writer:
|
await writer.send(Exception(err))
|
||||||
await writer.send(Exception(err))
|
|
||||||
return
|
return
|
||||||
|
|
||||||
async def _handle_get_request(self, request: Request, send: Send) -> None: # pragma: no cover
|
async def _handle_get_request(self, request: Request, send: Send) -> None: # pragma: no cover
|
||||||
|
|
@ -706,13 +707,15 @@ class StreamableHTTPServerTransport:
|
||||||
return
|
return
|
||||||
|
|
||||||
# Create SSE stream
|
# 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():
|
async def standalone_sse_writer():
|
||||||
try:
|
try:
|
||||||
# Create a standalone message stream for server-initiated messages
|
# 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]
|
standalone_stream_reader = self._request_streams[GET_STREAM_KEY][1]
|
||||||
|
|
||||||
async with sse_stream_writer, standalone_stream_reader:
|
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)
|
replay_protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
|
||||||
|
|
||||||
# Create SSE stream for replay
|
# 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():
|
async def replay_sender():
|
||||||
try:
|
try:
|
||||||
|
|
@ -918,22 +921,32 @@ class StreamableHTTPServerTransport:
|
||||||
|
|
||||||
# If stream ID not in mapping, create it
|
# If stream ID not in mapping, create it
|
||||||
if stream_id and stream_id not in self._request_streams:
|
if stream_id and stream_id not in self._request_streams:
|
||||||
# Register SSE writer so close_sse_stream() can close it
|
try:
|
||||||
self._sse_stream_writers[stream_id] = sse_stream_writer
|
# 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
|
# Prime the resumed connection so the client sees the stream
|
||||||
await self._maybe_send_priming_event(stream_id, sse_stream_writer, replay_protocol_version)
|
# 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
|
# Create new request streams for this connection
|
||||||
self._request_streams[stream_id] = anyio.create_memory_object_stream[EventMessage](0)
|
self._request_streams[stream_id] = anyio.create_memory_object_stream[EventMessage](
|
||||||
msg_reader = self._request_streams[stream_id][1]
|
REQUEST_STREAM_BUFFER_SIZE
|
||||||
|
)
|
||||||
|
msg_reader = self._request_streams[stream_id][1]
|
||||||
|
|
||||||
# Forward messages to SSE
|
# Forward messages to SSE
|
||||||
async with msg_reader:
|
async with msg_reader:
|
||||||
async for event_message in msg_reader:
|
async for event_message in msg_reader:
|
||||||
event_data = self._create_event_data(event_message)
|
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:
|
except anyio.ClosedResourceError:
|
||||||
# Expected when close_sse_stream() is called
|
# Expected when close_sse_stream() is called
|
||||||
logger.debug("Replay SSE stream closed by close_sse_stream()")
|
logger.debug("Replay SSE stream closed by close_sse_stream()")
|
||||||
|
|
|
||||||
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
|
@pytest.mark.anyio
|
||||||
async def test_priming_event_not_sent_for_old_protocol_version():
|
async def test_priming_event_not_minted_for_old_protocol_version():
|
||||||
"""Test that _maybe_send_priming_event skips for old protocol versions (backwards compat)."""
|
"""`_mint_priming_event` returns None for pre-2025-11-25 clients (backwards compat)."""
|
||||||
# Create a transport with an event store
|
|
||||||
transport = StreamableHTTPServerTransport(
|
transport = StreamableHTTPServerTransport(
|
||||||
"/mcp",
|
"/mcp",
|
||||||
event_store=SimpleEventStore(),
|
event_store=SimpleEventStore(),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Create a mock stream writer
|
assert await transport._mint_priming_event("test-request-id", "2025-06-18") is None
|
||||||
write_stream, read_stream = anyio.create_memory_object_stream[dict[str, Any]](1)
|
event = await transport._mint_priming_event("test-request-id-2", "2025-11-25")
|
||||||
|
assert event is not None
|
||||||
try:
|
assert event["data"] == ""
|
||||||
# Call _maybe_send_priming_event with OLD protocol version - should NOT send
|
assert "retry" not in event
|
||||||
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
|
@pytest.mark.anyio
|
||||||
async def test_priming_event_not_sent_without_event_store():
|
async def test_priming_event_not_minted_without_event_store():
|
||||||
"""Test that _maybe_send_priming_event returns early when no event_store is configured."""
|
"""`_mint_priming_event` returns None when no event store is configured."""
|
||||||
# Create a transport WITHOUT an event store
|
|
||||||
transport = StreamableHTTPServerTransport("/mcp")
|
transport = StreamableHTTPServerTransport("/mcp")
|
||||||
|
|
||||||
# Create a mock stream writer
|
assert await transport._mint_priming_event("test-request-id", "2025-11-25") is None
|
||||||
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
|
@pytest.mark.anyio
|
||||||
async def test_priming_event_includes_retry_interval():
|
async def test_priming_event_includes_retry_interval():
|
||||||
"""Test that _maybe_send_priming_event includes retry field when retry_interval is set."""
|
"""`_mint_priming_event` carries the configured `retry` field."""
|
||||||
# Create a transport with an event store AND retry_interval
|
|
||||||
transport = StreamableHTTPServerTransport(
|
transport = StreamableHTTPServerTransport(
|
||||||
"/mcp",
|
"/mcp",
|
||||||
event_store=SimpleEventStore(),
|
event_store=SimpleEventStore(),
|
||||||
retry_interval=5000,
|
retry_interval=5000,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Create a mock stream writer
|
event = await transport._mint_priming_event("test-request-id", "2025-11-25")
|
||||||
write_stream, read_stream = anyio.create_memory_object_stream[dict[str, Any]](1)
|
assert event is not None
|
||||||
|
assert event["retry"] == 5000
|
||||||
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
|
@pytest.mark.anyio
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue