Make client-side cancellation work over the 2026 transports (#3046)

tests/shared/test_dispatcher.py

@@ -19,6 +19,7 @@ from mcp_types import (
     INVALID_REQUEST,
     REQUEST_TIMEOUT,
     ErrorData,
+    RequestId,
     Tool,
 )
 
@@ -396,6 +397,118 @@ async def test_direct_close_makes_run_return():
             server.close()
 
 
+@pytest.mark.anyio
+async def test_send_raw_request_honors_caller_supplied_request_id_verbatim_typed(pair_factory: PairFactory):
+    """A caller-supplied `CallOptions["request_id"]` reaches the peer's context verbatim —
+    "7" stays a string, never the integer 7 — and the next call without one still mints
+    a dispatcher id as before."""
+    async with running_pair(pair_factory) as (client, _server, _crec, srec):
+        with anyio.fail_after(5):
+            await client.send_raw_request("first", None, {"request_id": "7"})
+            await client.send_raw_request("second", None)
+    supplied, minted = (ctx.request_id for ctx in srec.contexts)
+    assert supplied == "7"
+    assert type(supplied) is str
+    assert type(minted) is int
+
+
+@pytest.mark.anyio
+async def test_send_raw_request_with_in_flight_request_id_raises_and_frees_id_on_completion(
+    pair_factory: PairFactory,
+):
+    """Reusing an id while it is in flight is a loud `ValueError` — silent reuse would
+    corrupt response correlation. Once the first request completes, the id is free
+    again: the reservation is in-flight-scoped, not permanent."""
+    entered = anyio.Event()
+    release = anyio.Event()
+
+    async def parked(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        entered.set()
+        await release.wait()
+        return {"served": method}
+
+    async with running_pair(pair_factory, server_on_request=parked) as (client, *_):
+        with anyio.fail_after(5):
+            async with anyio.create_task_group() as tg:
+
+                async def first() -> None:
+                    await client.send_raw_request("slow", None, {"request_id": "listen-1"})
+
+                tg.start_soon(first)
+                await entered.wait()
+                with pytest.raises(ValueError, match="already in flight"):
+                    await client.send_raw_request("duplicate", None, {"request_id": "listen-1"})
+                release.set()
+            result = await client.send_raw_request("again", None, {"request_id": "listen-1"})
+    assert result == {"served": "again"}
+
+
+@pytest.mark.anyio
+async def test_minted_ids_skip_a_caller_supplied_id_still_in_flight(pair_factory: PairFactory):
+    """The dispatcher mints PAST a key a supplied id occupies — the collision error
+    is reserved for the caller who chose the id, never an innocent minted request."""
+    entered = anyio.Event()
+    release = anyio.Event()
+    seen_ids: list[RequestId | None] = []
+
+    async def maybe_park(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        seen_ids.append(ctx.request_id)
+        if method == "park":
+            entered.set()
+            await release.wait()
+        return {}
+
+    async with running_pair(pair_factory, server_on_request=maybe_park) as (client, *_):
+        with anyio.fail_after(5):
+            async with anyio.create_task_group() as tg:
+
+                async def parked() -> None:
+                    await client.send_raw_request("park", None, {"request_id": "3"})
+
+                tg.start_soon(parked)
+                await entered.wait()
+                # The counter mints 1 and 2, then skips the occupied 3 to 4.
+                for _ in range(3):
+                    await client.send_raw_request("plain", None)
+                release.set()
+            assert [request_id for request_id in seen_ids if request_id != "3"] == [1, 2, 4]
+
+
+@pytest.mark.anyio
+async def test_supplied_numeric_string_id_collides_with_its_int_twin(pair_factory: PairFactory):
+    """ "7" and 7 are one id in the collision domain on BOTH dispatchers, so the
+    in-memory pair raises exactly where the wire dispatcher (whose pending keys
+    are coerced for response correlation) would."""
+    entered = anyio.Event()
+    release = anyio.Event()
+
+    async def parked(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        entered.set()
+        await release.wait()
+        return {}
+
+    async with running_pair(pair_factory, server_on_request=parked) as (client, *_):
+        with anyio.fail_after(5):
+            async with anyio.create_task_group() as tg:
+
+                async def first() -> None:
+                    await client.send_raw_request("slow", None, {"request_id": 7})
+
+                tg.start_soon(first)
+                await entered.wait()
+                with pytest.raises(ValueError, match="already in flight"):
+                    await client.send_raw_request("duplicate", None, {"request_id": "7"})
+                release.set()
+            # Completion frees the id for either spelling.
+            assert await client.send_raw_request("again", None, {"request_id": "7"}) == {}
+
+
 if TYPE_CHECKING:
     _d: Dispatcher[TransportContext] = DirectDispatcher(TransportContext(kind="direct", can_send_request=True))
     _o: Outbound = _d

tests/shared/test_jsonrpc_dispatcher.py

@@ -34,11 +34,10 @@ from mcp import Client
 from mcp.server import Server, ServerRequestContext
 from mcp.shared._compat import resync_tracer
 from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream
-from mcp.shared.dispatcher import CallOptions, DispatchContext
+from mcp.shared.dispatcher import CallOptions, DispatchContext, coerce_request_id
 from mcp.shared.exceptions import MCPError, NoBackChannelError
 from mcp.shared.jsonrpc_dispatcher import (  # pyright: ignore[reportPrivateUsage]
     JSONRPCDispatcher,
-    _coerce_id,
     _OutboundPlan,
     _Pending,
     _plan_outbound,
@@ -1821,7 +1820,7 @@ async def test_response_with_string_id_correlates_to_int_keyed_pending_request()
 
 @pytest.mark.anyio
 async def test_error_response_with_string_id_correlates_to_int_keyed_pending_request():
-    """A JSONRPCError echoing the request ID as a JSON string still resolves the waiter (same `_coerce_id` path)."""
+    """A JSONRPCError echoing the request ID as a JSON string still resolves the waiter (`coerce_request_id` path)."""
     c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
     s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
     client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
@@ -1900,10 +1899,10 @@ async def test_progress_with_string_token_reaches_callback_for_int_keyed_request
     assert seen == [0.5]
 
 
-def test_coerce_id_passes_through_non_numeric_string_and_int():
-    assert _coerce_id("7") == 7
-    assert _coerce_id("not-an-int") == "not-an-int"
-    assert _coerce_id(42) == 42
+def test_coerce_request_id_passes_through_non_numeric_string_and_int():
+    assert coerce_request_id("7") == 7
+    assert coerce_request_id("not-an-int") == "not-an-int"
+    assert coerce_request_id(42) == 42
 
 
 @pytest.mark.anyio
@@ -2154,7 +2153,7 @@ async def test_request_with_bool_meta_progress_token_is_not_adopted():
     ids=["string-cancel-for-int-request", "int-cancel-for-string-request"],
 )
 async def test_cancelled_correlates_across_string_and_int_request_id_forms(request_id: RequestId, cancel_id: object):
-    """A peer that stringifies the id between request and cancel still cancels (same `_coerce_id` path)."""
+    """A peer that stringifies the id between request and cancel still cancels (same `coerce_request_id` path)."""
     c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
     s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
     server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
@@ -2381,3 +2380,38 @@ async def test_server_middleware_observes_cancelled_notification():
     assert observed[0][0] == "notifications/cancelled"
     assert observed[0][1]["requestId"] == request_id
     assert observed[0][1]["reason"] == "user clicked stop"
+
+
+@pytest.mark.anyio
+async def test_send_raw_request_with_caller_supplied_string_id_is_verbatim_on_the_wire():
+    """A supplied "7" goes on the wire as the string "7", and the response still
+    correlates when the peer echoes it back as the integer 7 — the pending key gets
+    the same coercion `_resolve_pending` applies to inbound ids."""
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
+    client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
+    on_request, on_notify = echo_handlers(Recorder())
+    result_box: list[dict[str, Any]] = []
+    done = anyio.Event()
+    try:
+        with anyio.fail_after(5):
+            async with anyio.create_task_group() as tg:
+
+                async def call() -> None:
+                    result_box.append(await client.send_raw_request("tools/list", None, {"request_id": "7"}))
+                    done.set()
+
+                await tg.start(client.run, on_request, on_notify)
+                tg.start_soon(call)
+                wire = await c2s_recv.receive()
+                assert isinstance(wire, SessionMessage)
+                assert isinstance(wire.message, JSONRPCRequest)
+                assert wire.message.id == "7"
+                assert type(wire.message.id) is str
+                await s2c_send.send(SessionMessage(JSONRPCResponse(jsonrpc="2.0", id=7, result={"ok": True})))
+                await done.wait()
+                tg.cancel_scope.cancel()
+    finally:
+        for stream in (c2s_send, c2s_recv, s2c_send, s2c_recv):
+            stream.close()
+    assert result_box == [{"ok": True}]