Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion src/mcp/shared/direct_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
9 changes: 8 additions & 1 deletion src/mcp/shared/jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {})
Expand Down
31 changes: 31 additions & 0 deletions tests/shared/test_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading