fix(shared): stop concurrent int/string request ids from colliding in _in_flight

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

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."""