feat: JSONRPCDispatcher receive loop and dispatch (chunk b)

maxisbey · maxisbey · commit 136a22d9b728 · 2026-04-16T21:43:14.000Z
run() drives the receive loop in a per-request task group;
task_status.started() fires once send_request is usable. _dispatch routes each
inbound message synchronously (no awaits — send_nowait/_spawn only) to avoid
head-of-line blocking. _spawn propagates the sender's contextvars via
Context.run(tg.start_soon, ...) so auth/OTel set by ASGI middleware survive.
_fan_out_closed wakes pending send_request waiters with CONNECTION_CLOSED on
shutdown (called both post-EOF and in finally; idempotent).

Wire-param extraction (progressToken, cancelled.requestId, progress fields)
uses structural match patterns — runtime narrowing, no casts, no mcp.types
model coupling; malformed input fails to match and the correlation is skipped.

_handle_request is happy-path only here (run on_request, write response); the
exception-to-wire boundary lands in the next commit.

Dispatcher.run() Protocol gained a task_status kwarg (it's a contract-level
guarantee). DirectDispatcher.run() updated to match. running_pair now uses
tg.start so the test body runs only once the dispatcher is ready.

20 contract tests pass; the 2 needing the exception boundary are strict-xfail.
diff --git a/src/mcp/shared/direct_dispatcher.py b/src/mcp/shared/direct_dispatcher.py
@@ -20,6 +20,7 @@
 from typing import Any
 
 import anyio
+import anyio.abc
 
 from mcp.shared.dispatcher import CallOptions, OnNotify, OnRequest, ProgressFnT
 from mcp.shared.exceptions import MCPError, NoBackChannelError
@@ -101,10 +102,17 @@ async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
             raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
         await self._peer._dispatch_notify(method, params)
 
-    async def run(self, on_request: OnRequest, on_notify: OnNotify) -> None:
+    async def run(
+        self,
+        on_request: OnRequest,
+        on_notify: OnNotify,
+        *,
+        task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
+    ) -> None:
         self._on_request = on_request
         self._on_notify = on_notify
         self._ready.set()
+        task_status.started()
         await self._closed.wait()
 
     def close(self) -> None:
diff --git a/src/mcp/shared/dispatcher.py b/src/mcp/shared/dispatcher.py
@@ -20,6 +20,7 @@
 from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable
 
 import anyio
+import anyio.abc
 
 from mcp.shared.transport_context import TransportContext
 
@@ -136,11 +137,21 @@ class Dispatcher(Outbound, Protocol[TransportT_co]):
     receive loop, per-request concurrency, and cancellation/progress wiring.
     """
 
-    async def run(self, on_request: OnRequest, on_notify: OnNotify) -> None:
+    async def run(
+        self,
+        on_request: OnRequest,
+        on_notify: OnNotify,
+        *,
+        task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
+    ) -> None:
         """Drive the receive loop until the underlying channel closes.
 
         Each inbound request is dispatched to ``on_request`` in its own task;
         the returned dict (or raised ``MCPError``) is sent back as the response.
         Inbound notifications go to ``on_notify``.
+
+        ``task_status.started()`` is called once the dispatcher is ready to
+        accept ``send_request``/``notify`` calls, so callers can use
+        ``await tg.start(dispatcher.run, on_request, on_notify)``.
         """
         ...
diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py
@@ -8,20 +8,30 @@
 The MCP type layer (`ServerRunner`, `Context`, `Client`) sits above this and
 sees only `(ctx, method, params) -> dict`. Transports sit below and see only
 `SessionMessage` reads/writes.
+
+The dispatcher is *mostly* MCP-agnostic — methods/params are opaque strings and
+dicts — but it intercepts ``notifications/cancelled`` and
+``notifications/progress`` because request correlation, cancellation and
+progress are exactly the wiring this layer exists to provide. Those few wire
+shapes are extracted with structural ``match`` patterns (no casts, no
+``mcp.types`` model coupling); a malformed payload simply fails to match and
+the correlation is skipped.
 """
 
 from __future__ import annotations
 
+import contextvars
 import logging
-from collections.abc import Callable, Mapping
+from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass, field
-from typing import Any, Generic, Literal, TypeVar, overload
+from typing import Any, Generic, Literal, TypeVar, cast, overload
 
 import anyio
+import anyio.abc
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
 from mcp.shared._stream_protocols import ReadStream, WriteStream
-from mcp.shared.dispatcher import CallOptions, OnNotify, OnRequest, ProgressFnT
+from mcp.shared.dispatcher import CallOptions, Dispatcher, OnNotify, OnRequest, ProgressFnT
 from mcp.shared.exceptions import MCPError, NoBackChannelError
 from mcp.shared.message import (
     ClientMessageMetadata,
@@ -31,11 +41,14 @@
 )
 from mcp.shared.transport_context import TransportContext
 from mcp.types import (
+    CONNECTION_CLOSED,
     REQUEST_TIMEOUT,
     ErrorData,
+    JSONRPCError,
     JSONRPCMessage,
     JSONRPCNotification,
     JSONRPCRequest,
+    JSONRPCResponse,
     ProgressToken,
     RequestId,
 )
@@ -141,8 +154,12 @@ def _outbound_metadata(related_request_id: RequestId | None, opts: CallOptions |
     return None
 
 
-class JSONRPCDispatcher(Generic[TransportT]):
-    """`Dispatcher` over the existing `SessionMessage` stream contract."""
+class JSONRPCDispatcher(Dispatcher[TransportT]):
+    """`Dispatcher` over the existing `SessionMessage` stream contract.
+
+    Inherits the `Dispatcher` Protocol explicitly so pyright checks
+    conformance at the class definition rather than at first use.
+    """
 
     @overload
     def __init__(
@@ -171,13 +188,20 @@ def __init__(
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
-        self._transport_builder = transport_builder or _default_transport_builder
+        # The overloads guarantee that when `transport_builder` is omitted,
+        # `TransportT` is `TransportContext`, so the default is type-correct;
+        # pyright can't see across overloads, hence the cast.
+        self._transport_builder = cast(
+            "Callable[[RequestId | None, MessageMetadata], TransportT]",
+            transport_builder or _default_transport_builder,
+        )
         self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode
         self._raise_handler_exceptions = raise_handler_exceptions
 
         self._next_id = 0
         self._pending: dict[RequestId, _Pending] = {}
         self._in_flight: dict[RequestId, _InFlight[TransportT]] = {}
+        self._tg: anyio.abc.TaskGroup | None = None
         self._running = False
 
     async def send_request(
@@ -219,6 +243,11 @@ async def send_request(
             meta["progressToken"] = request_id
             out_params = {**(out_params or {}), "_meta": meta}
 
+        # buffer=1: at most one outcome is ever delivered. A `WouldBlock` from
+        # `_resolve_pending`/`_fan_out_closed` means the waiter already has an
+        # outcome and dropping the late/redundant signal is correct. buffer=0
+        # is unsafe — there's a window between registering `_pending[id]` and
+        # parking in `receive()` where a close signal would be lost.
         send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
         pending = _Pending(send=send, receive=receive, on_progress=on_progress)
         self._pending[request_id] = pending
@@ -264,8 +293,197 @@ async def notify(
         msg = JSONRPCNotification(jsonrpc="2.0", method=method, params=dict(params) if params is not None else None)
         await self._write(msg, _outbound_metadata(_related_request_id, None))
 
-    async def run(self, on_request: OnRequest, on_notify: OnNotify) -> None:
-        raise NotImplementedError  # chunk (b)
+    async def run(
+        self,
+        on_request: OnRequest,
+        on_notify: OnNotify,
+        *,
+        task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
+    ) -> None:
+        """Drive the receive loop until the read stream closes.
+
+        Each inbound request is handled in its own task in an internal task
+        group; ``task_status.started()`` fires once that group is open, so
+        ``await tg.start(dispatcher.run, ...)`` resumes when ``send_request``
+        is usable.
+        """
+        try:
+            async with anyio.create_task_group() as tg:
+                self._tg = tg
+                self._running = True
+                task_status.started()
+                async with self._read_stream:
+                    async for item in self._read_stream:
+                        # Duck-typed: `_context_streams.ContextReceiveStream`
+                        # exposes `.last_context` (the sender's contextvars
+                        # snapshot per message). Plain memory streams don't.
+                        sender_ctx: contextvars.Context | None = getattr(self._read_stream, "last_context", None)
+                        self._dispatch(item, on_request, on_notify, sender_ctx)
+                # Read stream EOF: wake any blocked `send_request` waiters now,
+                # *before* the task group joins, so handlers parked in
+                # `dctx.send_request()` can unwind and the join doesn't deadlock.
+                self._running = False
+                self._fan_out_closed()
+        finally:
+            # Covers the cancel/crash paths where the inline fan-out above is
+            # never reached. Idempotent.
+            self._running = False
+            self._tg = None
+            self._fan_out_closed()
+
+    def _dispatch(
+        self,
+        item: SessionMessage | Exception,
+        on_request: OnRequest,
+        on_notify: OnNotify,
+        sender_ctx: contextvars.Context | None,
+    ) -> None:
+        """Route one inbound item. Synchronous: never awaits.
+
+        Everything here is `send_nowait` or `_spawn`. An `await` would let one
+        slow message head-of-line block the entire read loop.
+        """
+        if isinstance(item, Exception):
+            logger.debug("transport yielded exception: %r", item)
+            return
+        metadata = item.metadata
+        msg = item.message
+        match msg:
+            case JSONRPCRequest():
+                self._dispatch_request(msg, metadata, on_request, sender_ctx)
+            case JSONRPCNotification():
+                self._dispatch_notification(msg, metadata, on_notify, sender_ctx)
+            case JSONRPCResponse():
+                self._resolve_pending(msg.id, msg.result)
+            case JSONRPCError():
+                # `id` may be None per JSON-RPC (parse error before id known).
+                self._resolve_pending(msg.id, msg.error)
+
+    def _dispatch_request(
+        self,
+        req: JSONRPCRequest,
+        metadata: MessageMetadata,
+        on_request: OnRequest,
+        sender_ctx: contextvars.Context | None,
+    ) -> None:
+        progress_token: ProgressToken | None
+        match req.params:
+            case {"_meta": {"progressToken": str() | int() as progress_token}}:
+                pass
+            case _:
+                progress_token = None
+        transport_ctx = self._transport_builder(req.id, metadata)
+        dctx = _JSONRPCDispatchContext(
+            transport=transport_ctx,
+            _dispatcher=self,
+            _request_id=req.id,
+            _progress_token=progress_token,
+        )
+        scope = anyio.CancelScope()
+        self._in_flight[req.id] = _InFlight(scope=scope, dctx=dctx)
+        self._spawn(self._handle_request, req, dctx, scope, on_request, sender_ctx=sender_ctx)
+
+    def _dispatch_notification(
+        self,
+        msg: JSONRPCNotification,
+        metadata: MessageMetadata,
+        on_notify: OnNotify,
+        sender_ctx: contextvars.Context | None,
+    ) -> None:
+        if msg.method == "notifications/cancelled":
+            match msg.params:
+                case {"requestId": str() | int() as rid} if (in_flight := self._in_flight.get(rid)) is not None:
+                    in_flight.cancelled_by_peer = True
+                    in_flight.dctx.cancel_requested.set()
+                    if self._peer_cancel_mode == "interrupt":
+                        in_flight.scope.cancel()
+                case _:
+                    pass
+            return
+        if msg.method == "notifications/progress":
+            match msg.params:
+                case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if (
+                    pending := self._pending.get(token)
+                ) is not None and pending.on_progress is not None:
+                    total = msg.params.get("total")
+                    message = msg.params.get("message")
+                    self._spawn(
+                        pending.on_progress,
+                        float(progress),
+                        float(total) if isinstance(total, int | float) else None,
+                        message if isinstance(message, str) else None,
+                        sender_ctx=sender_ctx,
+                    )
+                case _:
+                    pass
+            # fall through: progress is also teed to on_notify
+        transport_ctx = self._transport_builder(None, metadata)
+        dctx = _JSONRPCDispatchContext(transport=transport_ctx, _dispatcher=self, _request_id=None)
+        self._spawn(on_notify, dctx, msg.method, msg.params, sender_ctx=sender_ctx)
+
+    def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None:
+        pending = self._pending.get(request_id) if request_id is not None else None
+        if pending is None:
+            logger.debug("dropping response for unknown/late request id %r", request_id)
+            return
+        try:
+            pending.send.send_nowait(outcome)
+        except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError):
+            logger.debug("waiter for request id %r already gone", request_id)
+
+    def _spawn(
+        self,
+        fn: Callable[..., Awaitable[Any]],
+        *args: object,
+        sender_ctx: contextvars.Context | None,
+    ) -> None:
+        """Schedule ``fn(*args)`` in the run() task group, propagating the sender's contextvars.
+
+        ASGI middleware (auth, OTel) sets contextvars on the request task that
+        wrote into the read stream. ``Context.run(tg.start_soon, ...)`` makes
+        the spawned handler inherit *that* context instead of the receive
+        loop's, so ``auth_context_var`` and OTel spans survive.
+        """
+        assert self._tg is not None
+        if sender_ctx is not None:
+            sender_ctx.run(self._tg.start_soon, fn, *args)
+        else:
+            self._tg.start_soon(fn, *args)
+
+    def _fan_out_closed(self) -> None:
+        """Wake every pending ``send_request`` waiter with ``CONNECTION_CLOSED``.
+
+        Synchronous (uses ``send_nowait``) because it's called from ``finally``
+        which may be inside a cancelled scope. Idempotent.
+        """
+        closed = ErrorData(code=CONNECTION_CLOSED, message="connection closed")
+        for pending in self._pending.values():
+            try:
+                pending.send.send_nowait(closed)
+            except (anyio.WouldBlock, anyio.BrokenResourceError, anyio.ClosedResourceError):
+                pass
+        self._pending.clear()
+
+    async def _handle_request(
+        self,
+        req: JSONRPCRequest,
+        dctx: _JSONRPCDispatchContext[TransportT],
+        scope: anyio.CancelScope,
+        on_request: OnRequest,
+    ) -> None:
+        """Run ``on_request`` for one inbound request and write its response.
+
+        Chunk (b): happy-path only. The full exception-to-wire boundary
+        (MCPError, ValidationError, INTERNAL_ERROR scrubbing, peer-cancel
+        no-response) lands in chunk (c).
+        """
+        try:
+            with scope:
+                result = await on_request(dctx, req.method, req.params)
+                await self._write(JSONRPCResponse(jsonrpc="2.0", id=req.id, result=result))
+        finally:
+            self._in_flight.pop(req.id, None)
+            dctx.close()
 
     def _allocate_id(self) -> int:
         self._next_id += 1
diff --git a/tests/shared/conftest.py b/tests/shared/conftest.py
@@ -48,20 +48,26 @@ def close() -> None:
     return client, server, close
 
 
-_JSONRPC_XFAIL = pytest.mark.xfail(
-    strict=True,
-    reason="JSONRPCDispatcher.run() not yet implemented (PR2 chunks b/c)",
-)
-
-
 @pytest.fixture(
     params=[
         pytest.param(direct_pair, id="direct"),
-        pytest.param(jsonrpc_pair, id="jsonrpc", marks=_JSONRPC_XFAIL),
+        pytest.param(jsonrpc_pair, id="jsonrpc"),
     ]
 )
 def pair_factory(request: pytest.FixtureRequest) -> PairFactory:
     return request.param
 
 
-__all__ = ["PairFactory", "direct_pair", "jsonrpc_pair"]
+def xfail_jsonrpc_chunk_c(request: pytest.FixtureRequest, factory: PairFactory) -> None:
+    """Apply a strict xfail when running against the JSON-RPC dispatcher.
+
+    Use for contract tests that require `_handle_request`'s exception boundary
+    (PR2 chunk c). Remove once that lands.
+    """
+    if factory is jsonrpc_pair:
+        request.applymarker(
+            pytest.mark.xfail(strict=True, reason="needs JSONRPCDispatcher._handle_request exception boundary")
+        )
+
+
+__all__ = ["PairFactory", "direct_pair", "jsonrpc_pair", "xfail_jsonrpc_chunk_c"]
diff --git a/tests/shared/test_dispatcher.py b/tests/shared/test_dispatcher.py
