diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 223373da0f..e27bd381c4 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -603,9 +603,12 @@ async def _dispatch_request( ) scope = anyio.CancelScope() # TODO(maxisbey): duplicate ids blind-overwrite (v1/TS parity); revisit - # rejecting with INVALID_REQUEST. Key coerced so a stringified - # `notifications/cancelled` id still correlates. - self._in_flight[coerce_request_id(req.id)] = _InFlight(scope=scope, dctx=dctx) + # rejecting with INVALID_REQUEST. Keyed by the verbatim id: distinct + # wire representations ("7" vs 7) are distinct requests, and coercion + # here would let a notifications/cancelled for one cancel the other. + # Cancellation correlation still finds peer-echoed ids via the + # coerced fallback in _dispatch_notification. + self._in_flight[req.id] = _InFlight(scope=scope, dctx=dctx) if req.method in self._inline_methods: # Spawn so `sender_ctx` applies, but park the read loop until the # handler returns - that's the inline ordering guarantee. @@ -639,10 +642,21 @@ def _dispatch_notification( """ if msg.method == "notifications/cancelled": rid = cancelled_request_id_from_params(msg.params) - if rid is not None and (in_flight := self._in_flight.get(coerce_request_id(rid))) is not None: - in_flight.dctx.cancel_requested.set() - if self._peer_cancel_mode == "interrupt": - in_flight.scope.cancel() + if rid is not None: + # Exact id first: "7" and 7 are distinct concurrent requests. + # Fall back to coerced/stringified forms so a peer that echoes + # a different type than the original request still cancels it. + entry = self._in_flight.get(rid) + if entry is None: + coerced = coerce_request_id(rid) + candidates = [coerced, str(coerced)] if coerced != rid else [str(rid)] + for candidate in candidates: + if candidate != rid and (entry := self._in_flight.get(candidate)) is not None: + break + if entry is not None: + entry.dctx.cancel_requested.set() + if self._peer_cancel_mode == "interrupt": + entry.scope.cancel() elif msg.method == "notifications/progress": match msg.params: case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if ( @@ -742,10 +756,10 @@ async def _handle_request( # Close the back-channel and drop from `_in_flight`; no checkpoint # since handler return, so a peer cancel can't interleave. # Identity guard: don't evict a duplicate id's newer entry. + # Keyed by the verbatim id, matching the store above. dctx.close() - key = coerce_request_id(req.id) - if (entry := self._in_flight.get(key)) is not None and entry.dctx is dctx: - del self._in_flight[key] + if (entry := self._in_flight.get(req.id)) is not None and entry.dctx is dctx: + del self._in_flight[req.id] if not dctx.cancel_requested.is_set(): # A write interrupted by cancellation may still have delivered # (a memory-stream send can hand its item to the receiver and diff --git a/tests/shared/test_jsonrpc_dispatcher.py b/tests/shared/test_jsonrpc_dispatcher.py index 37ca546cc9..60b6ab2fc6 100644 --- a/tests/shared/test_jsonrpc_dispatcher.py +++ b/tests/shared/test_jsonrpc_dispatcher.py @@ -185,6 +185,79 @@ async def on_unanswered() -> None: """Everything on the wire after a cancelled request 1: only the uncancelled control request 2 is answered.""" +@pytest.mark.anyio +async def test_peer_cancel_does_not_cross_request_id_types(): + """Concurrent requests "7" and 7 are distinct: a cancel for one must not + cancel the other. _in_flight used to key both by coerce_request_id, which + folds "7" and 7 into one slot, so the second request blind-overwrote the + first and a cancel for either cancelled whichever was stored.""" + c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](10) + s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage](10) + + str_id = "9007199254740991" + int_id = 9007199254740991 + answered_ids: set[object] = set() + + try: + server = JSONRPCDispatcher(c2s_recv, s2c_send) + + gate_released = anyio.Event() + both_started = anyio.Event() + started = 0 + + async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]: + nonlocal started + started += 1 + if started == 2: + both_started.set() + await gate_released.wait() + tag = params["arguments"]["tag"] if params and "arguments" in params else "" + return {"tag": tag} + + async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None: + pass + + async def on_unanswered() -> None: + pass + + async def drive() -> None: + for rid in (str_id, int_id): + msg = JSONRPCRequest( + jsonrpc="2.0", id=rid, method="tools/call", + params={"name": "gate", "arguments": {"tag": str(rid)}}, + ) + await c2s_send.send(SessionMessage(message=msg, metadata=ServerMessageMetadata(on_request_unanswered=on_unanswered))) + await both_started.wait() + # Cancel only the string-id request. + cancel = JSONRPCNotification(jsonrpc="2.0", method="notifications/cancelled", params={"requestId": str_id}) + await c2s_send.send(SessionMessage(message=cancel)) + # Release the gate so the int-id request can finish. + gate_released.set() + + async with anyio.create_task_group() as tg: + await tg.start(server.run, on_request, on_notify) + tg.start_soon(drive) + with anyio.fail_after(5): + while True: + msg = await s2c_recv.receive() + m = msg.message + if isinstance(m, JSONRPCResponse): + answered_ids.add(m.id) + if int_id in answered_ids: + break + tg.cancel_scope.cancel() + finally: + c2s_send.close() + c2s_recv.close() + s2c_send.close() + s2c_recv.close() + + # The string-id request was cancelled: it must have no response. + assert str_id not in answered_ids + # The int-id request must have completed and answered. + assert int_id in answered_ids + + @pytest.mark.anyio async def test_peer_cancel_interrupt_mode_interrupts_handler_and_writes_no_response(): """Spec MUST NOT: a cancelled request is never answered - not even with an error - and settles unanswered."""