-
Notifications
You must be signed in to change notification settings - Fork 3.7k
fix(server): reject concurrent duplicate JSON-RPC request ids #3221
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
|
@@ -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. | ||
|
|
@@ -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) | ||
|
|
||
| # 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: | ||
|
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]( | ||
|
|
@@ -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", | ||
|
|
@@ -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) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 Prompt for AI agents |
||
| # Extract related_request_id from meta if it exists | ||
| elif ( | ||
| session_message.metadata is not None | ||
|
|
@@ -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 | ||
|
|
||
|
|
||
There was a problem hiding this comment.
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