Synthesize request ids in DirectDispatcher

Server-initiated sampling/elicitation/roots requests over a ClientSession
built with dispatcher=DirectDispatcher failed before the callback ran:
the session requires a populated request id and direct dispatch carried
none. DirectDispatcher now assigns per-instance monotonic ids to inbound
requests (notifications keep None, which is how middleware distinguishes
them). Adds a non-ping direct-dispatch test and bounds the indefinite
awaits in the existing dispatcher= tests.

Author: maxisbey

src/mcp/client/session.py

@@ -507,7 +507,7 @@ async def _on_request(
 
         response: types.ClientResult | types.ErrorData
         if isinstance(request, types.PingRequest):
-            # Answered without a context: direct dispatch carries no request id.
+            # Answered without a context: ping has no callback that would need one.
             response = types.EmptyResult()
         else:
             assert dctx.request_id is not None  # the callback-driving dispatchers always assign ids

src/mcp/shared/direct_dispatcher.py

@@ -50,7 +50,7 @@ class _DirectDispatchContext:
     _back_request: _Request
     _back_notify: _Notify
     request_id: RequestId | None = None
-    """Always `None`: direct dispatch has no wire-level request id."""
+    """A dispatcher-synthesized id for requests; `None` for notifications."""
     message_metadata: MessageMetadata = None  # TODO(maxisbey): remove for Context rework
     """Always `None`: in-memory dispatch attaches no transport metadata."""
     _on_progress: ProgressFnT | None = None
@@ -91,6 +91,7 @@ def __init__(self, transport_ctx: TransportContext):
         self._peer: DirectDispatcher | None = None
         self._on_request: OnRequest | None = None
         self._on_notify: OnNotify | None = None
+        self._next_id = 0
         self._ready = anyio.Event()
         self._closed = anyio.Event()
 
@@ -128,13 +129,16 @@ async def run(
     def close(self) -> None:
         self._closed.set()
 
-    def _make_context(self, on_progress: ProgressFnT | None = None) -> _DirectDispatchContext:
+    def _make_context(
+        self, on_progress: ProgressFnT | None = None, request_id: RequestId | None = None
+    ) -> _DirectDispatchContext:
         assert self._peer is not None
         peer = self._peer
         return _DirectDispatchContext(
             transport=self._transport_ctx,
             _back_request=lambda m, p, o: peer._dispatch_request(m, p, o),
             _back_notify=lambda m, p: peer._dispatch_notify(m, p),
+            request_id=request_id,
             _on_progress=on_progress,
         )
 
@@ -147,7 +151,9 @@ async def _dispatch_request(
         await self._ready.wait()
         assert self._on_request is not None
         opts = opts or {}
-        dctx = self._make_context(on_progress=opts.get("on_progress"))
+        # Synthesize an id: the DispatchContext contract reserves None for notifications.
+        self._next_id += 1
+        dctx = self._make_context(on_progress=opts.get("on_progress"), request_id=self._next_id)
         try:
             with anyio.fail_after(opts.get("timeout")):
                 try:

tests/client/test_session.py

@@ -8,6 +8,7 @@
 import anyio.abc
 import anyio.streams.memory
 import pytest
+from pydantic import FileUrl
 
 from mcp import types
 from mcp.client import ClientRequestContext
@@ -972,18 +973,54 @@ async def server_on_notify(
 
     session = ClientSession(dispatcher=client_side)
     results: list[types.EmptyResult] = []
-    async with anyio.create_task_group() as tg:
-        await tg.start(server_side.run, server_on_request, server_on_notify)
-        async with session:
-            results.append(await session.send_ping(meta=None))
-            # Server-to-client: direct dispatch delivers ping with no params member (no _meta injection).
-            assert await server_side.send_raw_request("ping", None) == {}
-            await session.send_notification(types.RootsListChangedNotification())
-        server_side.close()
+    with anyio.fail_after(5):
+        async with anyio.create_task_group() as tg:
+            await tg.start(server_side.run, server_on_request, server_on_notify)
+            async with session:
+                results.append(await session.send_ping(meta=None))
+                # Server-to-client: direct dispatch delivers ping with no params member (no _meta injection).
+                assert await server_side.send_raw_request("ping", None) == {}
+                await session.send_notification(types.RootsListChangedNotification())
+            server_side.close()
     assert results == [types.EmptyResult()]
     assert notified == ["notifications/roots/list_changed"]
 
 
+@pytest.mark.anyio
+async def test_direct_dispatch_roots_list_reaches_callback_with_synthesized_request_id():
+    """A server-initiated roots/list over dispatcher= reaches the registered callback and round-trips
+    the result; the callback context carries an int request_id (SDK-defined: DirectDispatcher
+    synthesizes ids)."""
+    client_side, server_side = create_direct_dispatcher_pair()
+    contexts: list[ClientRequestContext] = []
+
+    async def list_roots(context: ClientRequestContext) -> types.ListRootsResult:
+        contexts.append(context)
+        return types.ListRootsResult(roots=[types.Root(uri=FileUrl("file:///workspace"))])
+
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
+    ) -> dict[str, object]:
+        raise NotImplementedError
+
+    async def server_on_notify(
+        ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
+    ) -> None:
+        raise NotImplementedError
+
+    session = ClientSession(dispatcher=client_side, list_roots_callback=list_roots)
+    result: dict[str, Any] | None = None
+    with anyio.fail_after(5):
+        async with anyio.create_task_group() as tg:
+            await tg.start(server_side.run, server_on_request, server_on_notify)
+            async with session:
+                result = await server_side.send_raw_request("roots/list", None)
+            server_side.close()
+    assert result == {"roots": [{"uri": "file:///workspace"}]}
+    assert len(contexts) == 1
+    assert isinstance(contexts[0].request_id, int)
+
+
 @pytest.mark.anyio
 async def test_initialize_opts_out_of_cancel_on_abandon_while_other_requests_leave_it_unset():
     """`send_request` passes `cancel_on_abandon=False` for `initialize` — the spec forbids
@@ -1021,9 +1058,10 @@ async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
             pass
 
     dispatcher = RecordingDispatcher()
-    async with ClientSession(dispatcher=dispatcher) as session:
-        await session.initialize()
-        await session.send_ping()
+    with anyio.fail_after(5):
+        async with ClientSession(dispatcher=dispatcher) as session:
+            await session.initialize()
+            await session.send_ping()
     opts_by_method = dict(dispatcher.calls)
     assert opts_by_method["initialize"].get("cancel_on_abandon") is False
     assert "cancel_on_abandon" not in opts_by_method["ping"]

tests/shared/test_dispatcher.py

@@ -228,13 +228,15 @@ async def test_ctx_message_metadata_is_none_when_transport_attaches_nothing(pair
 
 @pytest.mark.anyio
 async def test_ctx_request_id_exposes_inbound_id(pair_factory: PairFactory):
-    """JSON-RPC carries the wire id through; direct dispatch has none."""
+    """Every dispatcher assigns each inbound request a distinct int id; JSON-RPC carries
+    the wire id through, DirectDispatcher synthesizes one (SDK-defined)."""
     async with running_pair(pair_factory) as (client, _server, _crec, srec):
         with anyio.fail_after(5):
             await client.send_raw_request("tools/call", None)
             await client.send_raw_request("tools/call", None)
     a, b = (ctx.request_id for ctx in srec.contexts)
-    assert (a is None and b is None) or (isinstance(a, int) and isinstance(b, int) and a != b)
+    assert isinstance(a, int) and isinstance(b, int)
+    assert a != b
 
 
 @pytest.mark.anyio