diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py index 1a4e9939a4..46272ca07b 100644 --- a/src/mcp/server/streamable_http.py +++ b/src/mcp/server/streamable_http.py @@ -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: + 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) # 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 diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index aeef25a278..7fde7d41e6 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -93,6 +93,14 @@ def first_sse_data(response: httpx2.Response) -> dict[str, Any]: raise ValueError("No data event in SSE response") # pragma: no cover +async def next_sse_data(lines: AsyncIterator[str]) -> dict[str, Any]: + """Return the next SSE `data:` payload from a live line iterator, parsed as JSON.""" + while True: + line = await anext(lines) + if line.startswith("data: "): + return json.loads(line.removeprefix("data: ")) + + def extract_protocol_version_from_sse(response: httpx2.Response) -> str: """Extract the negotiated protocol version from an SSE initialization response.""" return first_sse_data(response)["result"]["protocolVersion"] @@ -808,6 +816,443 @@ async def test_get_sse_stream(basic_app: Starlette) -> None: assert second_get.status_code == 409 +@pytest.mark.anyio +async def test_post_duplicate_request_id_rejected_while_first_still_in_flight(basic_app: Starlette) -> None: + """A second POST reusing an in-flight request's JSON-RPC id is rejected with 409; the first still completes. + + Per-request routing is keyed by request id, so overwriting an in-flight slot would + cross-wire responses (see #3137). Raw HTTP is required to assert the 409 status code. + """ + async with make_client(basic_app) as client: + init_response = await client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: init_response.headers[MCP_SESSION_ID_HEADER], + MCP_PROTOCOL_VERSION_HEADER: extract_protocol_version_from_sse(init_response), + } + shared_id = "shared-request-id" + + async with client.stream( + "POST", + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "wait_for_lock_with_notification", "arguments": {}}, + "id": shared_id, + }, + ) as first_response: + assert first_response.status_code == 200 + lines = first_response.aiter_lines() + # First notification confirms the request is registered and blocked on the lock. + with anyio.fail_after(5): + notification = await next_sse_data(lines) + assert notification["params"]["data"] == "First notification before lock" + + duplicate_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "test_tool", "arguments": {}}, + "id": shared_id, + }, + ) + assert duplicate_response.status_code == 409 + error = duplicate_response.json()["error"] + assert error["code"] == INVALID_REQUEST + assert "already in flight" in error["message"] + + release_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "release_lock", "arguments": {}}, + "id": "release-lock-1", + }, + ) + assert release_response.status_code == 200 + + with anyio.fail_after(5): + second_notification = await next_sse_data(lines) + final = await next_sse_data(lines) + assert second_notification["params"]["data"] == "Second notification after lock" + assert final["id"] == shared_id + assert final["result"]["content"][0]["text"] == "Completed" + + +@pytest.mark.anyio +async def test_json_response_duplicate_request_id_rejected_while_first_still_in_flight(json_app: Starlette) -> None: + """In JSON response mode, a second POST reusing an in-flight id is rejected with 409; the first still completes. + + Same collision as the SSE case, exercised on the JSON branch of `_handle_post_request` (#3137). + """ + async with make_client(json_app) as client: + init_response = await client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: init_response.headers[MCP_SESSION_ID_HEADER], + MCP_PROTOCOL_VERSION_HEADER: init_response.json()["result"]["protocolVersion"], + } + shared_id = "shared-json-request-id" + first_result: dict[str, Any] = {} + + async def run_first_request() -> None: + response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "wait_for_lock_with_notification", "arguments": {}}, + "id": shared_id, + }, + ) + assert response.status_code == 200 + first_result.update(response.json()) + + async with anyio.create_task_group() as tg: + tg.start_soon(run_first_request) + # First request registers its stream then blocks server-side on the lock; + # nothing else can run until that await yields. + await anyio.wait_all_tasks_blocked() + + duplicate_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "test_tool", "arguments": {}}, + "id": shared_id, + }, + ) + assert duplicate_response.status_code == 409 + error = duplicate_response.json()["error"] + assert error["code"] == INVALID_REQUEST + assert "already in flight" in error["message"] + + release_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "release_lock", "arguments": {}}, + "id": "release-lock-json-1", + }, + ) + assert release_response.status_code == 200 + + assert first_result["id"] == shared_id + assert first_result["result"]["content"][0]["text"] == "Completed" + + +@pytest.mark.anyio +async def test_request_id_reuse_after_completion_allowed(basic_app: Starlette) -> None: + """A request id may be reused once the earlier request with that id has completed. + + Only concurrent reuse is ambiguous to route; sequential reuse is allowed (some + clients send a constant id for every request). + """ + async with make_client(basic_app) as client: + init_response = await client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: init_response.headers[MCP_SESSION_ID_HEADER], + MCP_PROTOCOL_VERSION_HEADER: extract_protocol_version_from_sse(init_response), + } + + for _ in range(2): + response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "test_tool", "arguments": {}}, + }, + ) + assert response.status_code == 200 + body = first_sse_data(response) + assert body["id"] == 1 + assert body["result"]["content"][0]["text"] == "Called test_tool" + + +@pytest.mark.anyio +async def test_numeric_and_string_request_ids_do_not_collide_while_in_flight(basic_app: Starlette) -> None: + """Integer id 1 and string id "1" are distinct JSON-RPC ids and may run concurrently. + + Routing used to key streams with str(id), so an in-flight numeric 1 falsely + 409'd a string "1" (or cross-wired responses). Raw HTTP is required to control + the JSON-RPC id types on the wire. + """ + async with make_client(basic_app) as client: + init_response = await client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: init_response.headers[MCP_SESSION_ID_HEADER], + MCP_PROTOCOL_VERSION_HEADER: extract_protocol_version_from_sse(init_response), + } + + async with client.stream( + "POST", + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "wait_for_lock_with_notification", "arguments": {}}, + "id": 1, + }, + ) as first_response: + assert first_response.status_code == 200 + lines = first_response.aiter_lines() + with anyio.fail_after(5): + notification = await next_sse_data(lines) + assert notification["params"]["data"] == "First notification before lock" + + string_id_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "test_tool", "arguments": {}}, + "id": "1", + }, + ) + assert string_id_response.status_code == 200 + string_body = first_sse_data(string_id_response) + assert string_body["id"] == "1" + assert string_body["result"]["content"][0]["text"] == "Called test_tool" + + release_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "release_lock", "arguments": {}}, + "id": "release-lock-type-key", + }, + ) + assert release_response.status_code == 200 + + with anyio.fail_after(5): + second_notification = await next_sse_data(lines) + final = await next_sse_data(lines) + assert second_notification["params"]["data"] == "Second notification after lock" + assert final["id"] == 1 + assert final["result"]["content"][0]["text"] == "Completed" + + +@pytest.mark.anyio +async def test_request_id_matching_get_stream_key_does_not_block_standalone_get(basic_app: Starlette) -> None: + """A POST with string id \"_GET_stream\" must not occupy the standalone GET slot. + + Pre-fix, both used the bare string as the routing key, so a client could lock + itself out of GET (or have GET overwrite its POST slot). + """ + async with make_client(basic_app) as client: + init_response = await client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + session_id = init_response.headers[MCP_SESSION_ID_HEADER] + negotiated_version = extract_protocol_version_from_sse(init_response) + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, + } + + async with client.stream( + "POST", + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "wait_for_lock_with_notification", "arguments": {}}, + "id": "_GET_stream", + }, + ) as post_response: + assert post_response.status_code == 200 + lines = post_response.aiter_lines() + with anyio.fail_after(5): + notification = await next_sse_data(lines) + assert notification["params"]["data"] == "First notification before lock" + + get_headers = { + "Accept": "text/event-stream", + MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, + } + async with client.stream("GET", "/mcp", headers=get_headers) as get_response: + assert get_response.status_code == 200 + assert get_response.headers.get("Content-Type") == "text/event-stream" + + release_response = await client.post( + "/mcp", + headers=headers, + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "release_lock", "arguments": {}}, + "id": "release-after-get-key", + }, + ) + assert release_response.status_code == 200 + + with anyio.fail_after(5): + second_notification = await next_sse_data(lines) + final = await next_sse_data(lines) + assert second_notification["params"]["data"] == "Second notification after lock" + assert final["id"] == "_GET_stream" + assert final["result"]["content"][0]["text"] == "Completed" + + +@pytest.mark.anyio +async def test_duplicate_request_id_rejected_during_priming_event_store() -> None: + """The duplicate-id guard holds while the event store persists the priming event. + + With an event store, the SSE branch awaits `EventStore.store_event` before the + response starts. The routing slot is reserved before that await so a concurrent + POST reusing the id cannot slip past the guard and overwrite the slot (#3137). + """ + + class GatedEventStore(SimpleEventStore): + """Blocks the priming write for the gated stream until released.""" + + def __init__(self) -> None: + super().__init__() + self.entered = anyio.Event() + self.release = anyio.Event() + + async def store_event(self, stream_id: StreamId, message: types.JSONRPCMessage | None) -> EventId: + # Routing keys are type-tagged (`s:` for string JSON-RPC ids). + if stream_id == "s:gated-1" and message is None: + self.entered.set() + await self.release.wait() + return await super().store_event(stream_id, message) + + def first_message_sse_data(response: httpx2.Response) -> dict[str, Any]: + """First non-empty SSE data payload; skips the empty-data priming event.""" + for line in response.text.splitlines(): + if line.startswith("data: ") and line.removeprefix("data: ").strip(): + return json.loads(line.removeprefix("data: ")) + raise ValueError("No message data event in SSE response") # pragma: no cover + + store = GatedEventStore() + async with running_app(event_store=store) as app: + async with make_client(app) as client: + # Priming events require protocol >= 2025-11-25. + init_request = { + "jsonrpc": "2.0", + "method": "initialize", + "params": { + "clientInfo": {"name": "test-client", "version": "1.0"}, + "protocolVersion": types.LATEST_PROTOCOL_VERSION, + "capabilities": {}, + }, + "id": "init-1", + } + init_response = await client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=init_request, + ) + assert init_response.status_code == 200 + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: init_response.headers[MCP_SESSION_ID_HEADER], + MCP_PROTOCOL_VERSION_HEADER: first_message_sse_data(init_response)["result"]["protocolVersion"], + } + call: dict[str, Any] = { + "jsonrpc": "2.0", + "id": "gated-1", + "method": "tools/call", + "params": {"name": "test_tool", "arguments": {}}, + } + results: dict[str, httpx2.Response] = {} + + async def post_first() -> None: + results["first"] = await client.post("/mcp", headers=headers, json=call) + + async with anyio.create_task_group() as tg: + tg.start_soon(post_first) + with anyio.fail_after(5): + await store.entered.wait() + + # Without early slot reservation this POST would also suspend in the + # gated store and hang; fail_after turns that regression into a fast fail. + with anyio.fail_after(5): + duplicate = await client.post("/mcp", headers=headers, json=call) + assert duplicate.status_code == 409 + error = duplicate.json()["error"] + assert error["code"] == INVALID_REQUEST + assert "already in flight" in error["message"] + + store.release.set() + + assert results["first"].status_code == 200 + body = first_message_sse_data(results["first"]) + assert body["id"] == "gated-1" + assert body["result"]["content"][0]["text"] == "Called test_tool" + + @pytest.mark.anyio async def test_get_validation(basic_app: Starlette) -> None: """A GET without an Accept header covering text/event-stream is rejected with 406."""