Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 60 additions & 18 deletions src/mcp/server/streamable_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,32 @@
CONTENT_TYPE_JSON = "application/json"
CONTENT_TYPE_SSE = "text/event-stream"

# Special key for the standalone GET stream
# Type aliases
StreamId = str
EventId = str
# An SSE event-dict as accepted by sse-starlette (`event`, `data`, `id`, `retry`).
SSEEvent = dict[str, Any]

# Special key for the standalone GET stream. Request routing keys never equal
# this value (see `_request_stream_key`), so a JSON-RPC string id of
# `"_GET_stream"` cannot lock out or share the standalone GET slot.
GET_STREAM_KEY = "_GET_stream"


def _request_stream_key(request_id: RequestId) -> StreamId:
"""Type-preserving routing key for per-request streams and the event store.

JSON-RPC treats integer ``1`` and string ``\"1\"`` as distinct ids. Using
``str(id)`` collapses them (so an in-flight numeric id falsely 409s a
string id, or cross-wires responses) and also lets a string id equal to
``GET_STREAM_KEY`` share the standalone GET slot.
"""
# RequestId is int | str; exclude bool (a subclass of int) defensively.
if type(request_id) is int:
return f"i:{request_id}"
return f"s:{request_id}"


# 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.
Expand All @@ -76,12 +99,6 @@
# Pattern ensures entire string contains only valid characters by using ^ and $ anchors
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]


def check_accept_headers(request: Request) -> tuple[bool, bool]:
"""Return (has_json, has_sse) for the request's Accept header, with RFC 7231 wildcard handling.
Expand Down Expand Up @@ -598,7 +615,26 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
else request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
)

request_id = str(message.id)
# Type-preserving routing key so int 1 and str "1" stay distinct, and so
# a string id never collides with GET_STREAM_KEY. The membership test and
# the assignment below must not have an await between them: with
# resumability on, EventStore.store_event is user code and any await
# re-opens a check-then-act window (#3137 / #3163).
request_id = _request_stream_key(message.id)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: Direct callers of close_sse_stream() with its documented JSON-RPC request id now get a silent no-op because request streams are stored under encoded keys. Normalize this public method's input to the routing key while preserving the standalone GET key, and keep the internal callback aligned with that contract.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/mcp/server/streamable_http.py, line 623:

<comment>Direct callers of `close_sse_stream()` with its documented JSON-RPC request id now get a silent no-op because request streams are stored under encoded keys. Normalize this public method's input to the routing key while preserving the standalone GET key, and keep the internal callback aligned with that contract.</comment>

<file context>
@@ -598,16 +615,22 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
+            # the assignment below must not have an await between them: with
+            # resumability on, EventStore.store_event is user code and any await
+            # re-opens a check-then-act window (#3137 / #3163).
+            request_id = _request_stream_key(message.id)
+
+            # A second concurrent POST reusing an in-flight id would overwrite that
</file context>


# A second concurrent POST reusing an in-flight id would overwrite that
# slot and cross-wire responses (one caller receives the other's payload;
# the other hangs). Reject with 409 rather than silently re-routing. Ids
# may still be reused after the earlier request has completed; clients
# that retry after a timeout should send notifications/cancelled first.
if request_id in self._request_streams:
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
response = self._create_error_response(
f"Conflict: Request id {message.id!r} is already in flight for this "
f"session; send notifications/cancelled before reusing the id",
HTTPStatus.CONFLICT,
)
await response(scope, receive, send)
return

if self.is_json_response_enabled:
self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](
Expand Down Expand Up @@ -631,19 +667,25 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
await self._clean_up_memory_streams(request_id)
await response(scope, receive, send)
else:
# 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)

sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
self._sse_stream_writers[request_id] = sse_stream_writer
# Reserve the routing slot before any await so a concurrent POST
# reusing this id cannot pass the guard above while we mint the
# priming event (EventStore.store_event is user code and may
# await). Release the reservation if priming raises.
self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](
REQUEST_STREAM_BUFFER_SIZE
)
request_stream_reader = self._request_streams[request_id][1]

try:
# Still strictly precedes dispatch, so storage order == wire order.
priming_event = await self._mint_priming_event(request_id, protocol_version)
except BaseException:
await self._clean_up_memory_streams(request_id)
raise

sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
self._sse_stream_writers[request_id] = sse_stream_writer

headers = {
"Cache-Control": "no-cache, no-transform",
"Connection": "keep-alive",
Expand Down Expand Up @@ -1020,7 +1062,7 @@ async def message_router():
# Null-id errors (e.g., parse errors) fall through to
# the GET stream since they can't be correlated.
if isinstance(message, JSONRPCResponse | JSONRPCError) and message.id is not None:
target_request_id = str(message.id)
target_request_id = _request_stream_key(message.id)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3: The router regression test now times out because its manually registered streams use raw ids while responses route to s:-prefixed keys. Update its fixture keys to the same encoded representation (or route through the request registration path) so CI continues to cover the buffering behavior.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/mcp/server/streamable_http.py, line 1065:

<comment>The router regression test now times out because its manually registered streams use raw ids while responses route to `s:`-prefixed keys. Update its fixture keys to the same encoded representation (or route through the request registration path) so CI continues to cover the buffering behavior.</comment>

<file context>
@@ -1039,7 +1062,7 @@ async def message_router():
                         # the GET stream since they can't be correlated.
                         if isinstance(message, JSONRPCResponse | JSONRPCError) and message.id is not None:
-                            target_request_id = str(message.id)
+                            target_request_id = _request_stream_key(message.id)
                         # Extract related_request_id from meta if it exists
                         elif (
</file context>

# Extract related_request_id from meta if it exists
elif (
session_message.metadata is not None
Expand All @@ -1037,7 +1079,7 @@ async def message_router():
# storing or queueing rather than park it (#1764).
logger.debug(f"Dropped message related to request {related_request_id} in JSON mode")
continue
target_request_id = str(related_request_id)
target_request_id = _request_stream_key(related_request_id)

request_stream_id = target_request_id if target_request_id is not None else GET_STREAM_KEY

Expand Down
Loading
Loading