diff --git a/src/mcp/shared/direct_dispatcher.py b/src/mcp/shared/direct_dispatcher.py index e17283afa2..2442d4475f 100644 --- a/src/mcp/shared/direct_dispatcher.py +++ b/src/mcp/shared/direct_dispatcher.py @@ -117,6 +117,7 @@ def __init__(self, transport_ctx: TransportContext, *, raise_handler_exceptions: self._on_notify_intercept: OnNotifyIntercept | None = None self._next_id = 0 self._in_flight_ids: set[RequestId] = set() + self._retired_ids: set[int] = set() self._ready = anyio.Event() self._close_event = anyio.Event() self._running = False @@ -250,16 +251,21 @@ async def _dispatch_request( in_flight_key = coerce_request_id(request_id) if in_flight_key in self._in_flight_ids: raise ValueError(f"request id {request_id!r} is already in flight") + # Same no-reuse rule as JSONRPCDispatcher: retire the coerced + # key so a later minted request can't land on it. + if isinstance(in_flight_key, int): + self._retired_ids.add(in_flight_key) else: # Synthesize an id (the DispatchContext contract reserves None # for notifications), minting past any key a supplied id # occupies: the collision error is reserved for the caller # who actually chose the id. self._next_id += 1 - while self._next_id in self._in_flight_ids: + while self._next_id in self._in_flight_ids or self._next_id in self._retired_ids: self._next_id += 1 request_id = self._next_id in_flight_key = request_id + self._retired_ids = {key for key in self._retired_ids if key > request_id} self._in_flight_ids.add(in_flight_key) dctx = self._make_context(on_progress=opts.get("on_progress"), request_id=request_id) try: diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..16314f4b8f 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -307,6 +307,7 @@ def __init__( self._next_id = 0 self._pending: dict[RequestId, _Pending] = {} self._in_flight: dict[RequestId, _InFlight[TransportT]] = {} + self._retired_ids: set[int] = set() self._on_notify_intercept: OnNotifyIntercept | None = None self._tg: anyio.abc.TaskGroup | None = None self._running = False @@ -346,12 +347,18 @@ async def send_raw_request( pending_key = coerce_request_id(request_id) if pending_key in self._pending: raise ValueError(f"request id {request_id!r} is already in flight") + # Spec: an id is never reused in a session, even after completion — + # retire the coerced key so a later minted request can't land on it. + if isinstance(pending_key, int): + self._retired_ids.add(pending_key) else: # Mint past any key a supplied id occupies: the collision error is # reserved for the caller who actually chose the id. request_id = self._allocate_id() - while request_id in self._pending: + while request_id in self._pending or request_id in self._retired_ids: request_id = self._allocate_id() + # The counter never goes back, so retired keys below it are spent. + self._retired_ids = {key for key in self._retired_ids if key > request_id} pending_key = request_id out_params = dict(params) if params is not None else {} out_meta = dict(out_params.get("_meta") or {}) diff --git a/tests/shared/test_dispatcher.py b/tests/shared/test_dispatcher.py index c6ebb401ff..c07420fb7c 100644 --- a/tests/shared/test_dispatcher.py +++ b/tests/shared/test_dispatcher.py @@ -481,6 +481,37 @@ async def parked() -> None: assert [request_id for request_id in seen_ids if request_id != "3"] == [1, 2, 4] +@pytest.mark.anyio +async def test_minted_ids_advance_past_a_completed_caller_supplied_numeric_id(pair_factory: PairFactory): + """Spec: an id MUST NOT be reused by the requestor within a session — not even + after its request completed and left the in-flight set. Accepting a numeric + supplied id advances the mint counter past it, so minted ids never revisit it.""" + 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": 1}) + for _ in range(3): + await client.send_raw_request("plain", None) + supplied, *minted = (ctx.request_id for ctx in srec.contexts) + assert supplied == 1 + assert minted == [2, 3, 4] + + +@pytest.mark.anyio +async def test_minted_ids_advance_past_a_completed_supplied_numeric_string_id(pair_factory: PairFactory): + """The collision domain folds "7" and 7 into one key, so accepting the string form + retires 7 against future mints even though it left the in-flight set.""" + 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"}) + # Mints walk 1..9 but must skip the retired 7. + for _ in range(9): + await client.send_raw_request("plain", None) + supplied, *minted = (ctx.request_id for ctx in srec.contexts) + assert supplied == "7" + assert type(supplied) is str + assert minted == [1, 2, 3, 4, 5, 6, 8, 9, 10] + + @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