jsonrpc_dispatcher

class JSONRPCDispatcher(Dispatcher[TransportT]):
    """`Dispatcher` over the `SessionMessage` stream contract.

    Explicit Protocol base so pyright checks conformance at the class definition.
    """

    def __init__(
        self,
        read_stream: ReadStream[SessionMessage | Exception],
        write_stream: WriteStream[SessionMessage],
        *,
        transport_builder: Callable[[MessageMetadata], TransportT] | None = None,
        peer_cancel_mode: PeerCancelMode = "interrupt",
        raise_handler_exceptions: bool = False,
        inline_methods: frozenset[str] = frozenset(),
        on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None,
    ) -> None:
        """Wire a dispatcher over a transport's `SessionMessage` stream pair.

        Args:
            transport_builder: Builds each message's `TransportContext` from
                its `SessionMessage.metadata`.
            raise_handler_exceptions: Re-raise handler exceptions out of
                `run()` after the error response is written.
            inline_methods: Methods awaited in the read loop before the next
                message is dequeued (e.g. `initialize`); an inline handler
                that awaits the peer deadlocks the parked loop.
            on_stream_exception: Observer for `Exception` items on the read
                stream; without it they are debug-logged and dropped. Awaited
                inline in the read loop, so a slow observer stalls dispatch.
        """
        self._read_stream = read_stream
        self._write_stream = write_stream
        # With transport_builder omitted, TransportT defaults to
        # TransportContext; pyright can't connect the two, hence the cast.
        self._transport_builder = cast(
            "Callable[[MessageMetadata], TransportT]",
            transport_builder or _default_transport_builder,
        )
        self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode
        self._raise_handler_exceptions = raise_handler_exceptions
        self._inline_methods = inline_methods
        self.on_stream_exception = on_stream_exception
        """Observer for ``Exception`` items on the read stream. Mutable so a session can
        bind it after the dispatcher is built (e.g. ``ClientSession`` routing into
        ``message_handler``); only consulted inside ``run()`` so pre-enter assignment is safe."""

        # The correlation kernel owns the pending/in-flight tables; the
        # aliases keep the historical private names white-box tests read.
        self._corr: RequestCorrelator[_JSONRPCDispatchContext[TransportT]] = RequestCorrelator()
        self._pending: dict[RequestId, Pending] = self._corr.pending
        self._in_flight: dict[RequestId, InFlight[_JSONRPCDispatchContext[TransportT]]] = self._corr.in_flight
        self._on_notify_intercept: OnNotifyIntercept | None = None
        self._tg: anyio.abc.TaskGroup | None = None
        self._running = False

    async def send_raw_request(
        self,
        method: str,
        params: Mapping[str, Any] | None,
        opts: CallOptions | None = None,
        *,
        _related_request_id: RequestId | None = None,
    ) -> dict[str, Any]:
        """Send a JSON-RPC request and await its response.

        `_related_request_id` is set only by `_JSONRPCDispatchContext` so that
        mid-handler requests route onto the inbound request's SSE stream.

        Raises:
            MCPError: Peer error response; `REQUEST_TIMEOUT` if
                `opts["timeout"]` elapsed; `CONNECTION_CLOSED` if the
                transport closed or the dispatcher shut down.
            RuntimeError: Called before `run()`.
        """
        # Post-close sends get the same CONNECTION_CLOSED contract as in-flight
        # waiters (raised by the correlator); only a never-run dispatcher is a usage error.
        if not self._running and not self._corr.closed:
            raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run()")
        plan = _plan_outbound(_related_request_id, opts)
        return await self._corr.call(
            method,
            params,
            opts,
            write_request=partial(self._write, metadata=plan.metadata),
            send_cancel=partial(self._cancel_outbound, related_request_id=_related_request_id),
            cancel_on_abandon=plan.cancel_on_abandon,
        )

    async def notify(
        self,
        method: str,
        params: Mapping[str, Any] | None,
        opts: CallOptions | None = None,
        *,
        _related_request_id: RequestId | None = None,
    ) -> None:
        """Send a fire-and-forget notification.

        Fire-and-forget all the way: a post-close send or a write onto a
        torn-down transport drops the notification with a debug log instead
        of raising (same policy as the response writes and `ctx.notify`).
        """
        if self._corr.closed:
            logger.debug("dropped %s: dispatcher closed", method)
            return
        # Leave `params` unset when None: with `exclude_unset=True` an explicit
        # None would serialize as `"params": null`, which JSON-RPC 2.0 forbids.
        if params is not None:
            msg = JSONRPCNotification(jsonrpc="2.0", method=method, params=dict(params))
        else:
            msg = JSONRPCNotification(jsonrpc="2.0", method=method)
        try:
            await self._write(msg, _plan_outbound(_related_request_id, opts).metadata)
        except (anyio.BrokenResourceError, anyio.ClosedResourceError):
            # Transport tore down before run() noticed EOF.
            logger.debug("dropped %s: write stream closed", method)

    async def run(
        self,
        on_request: OnRequest,
        on_notify: OnNotify,
        on_notify_intercept: OnNotifyIntercept | None = None,
        *,
        task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
    ) -> None:
        """Drive the receive loop until the read stream closes.

        `task_status.started()` fires once `send_raw_request` is usable.
        Single-shot: once the loop ends the dispatcher stays closed and cannot be restarted.
        """
        self._on_notify_intercept = on_notify_intercept
        try:
            # LIFO exits: the write stream closes only after the task-group join, so teardown writes still land.
            async with self._write_stream:
                async with anyio.create_task_group() as tg:
                    self._tg = tg
                    self._running = True
                    task_status.started()
                    try:
                        async with self._read_stream:
                            try:
                                async for item in self._read_stream:
                                    # Duck-typed: only `ContextReceiveStream` carries the
                                    # sender's per-message contextvars snapshot.
                                    sender_ctx: contextvars.Context | None = getattr(
                                        self._read_stream, "last_context", None
                                    )
                                    await self._dispatch(item, on_request, on_notify, sender_ctx)
                            except anyio.ClosedResourceError:
                                # Receive end closed under us (stateless SHTTP teardown); same as EOF.
                                logger.debug("read stream closed by transport; treating as EOF")
                        # EOF: wake blocked `send_raw_request` waiters with CONNECTION_CLOSED.
                        self._running = False
                        self._corr.close()
                    finally:
                        # Cancel in-flight handlers; otherwise the task-group join
                        # waits on handlers whose callers are already gone.
                        tg.cancel_scope.cancel()
        finally:
            # Covers cancel/crash paths that skip the inline close; idempotent.
            self._running = False
            self._tg = None
            self._corr.close()
            await resync_tracer()

    async def _dispatch(
        self,
        item: SessionMessage | Exception,
        on_request: OnRequest,
        on_notify: OnNotify,
        sender_ctx: contextvars.Context | None,
    ) -> None:
        """Route one inbound item.

        Only `inline_methods` requests and the `on_stream_exception` observer
        are awaited; any other `await` would head-of-line block the read loop.
        """
        if isinstance(item, Exception):
            if self.on_stream_exception is None:
                logger.debug("transport yielded exception: %r", item)
                return
            try:
                await self.on_stream_exception(item)
            except Exception:
                logger.exception("on_stream_exception observer raised")
            return
        metadata = item.metadata
        msg = item.message
        match msg:
            case JSONRPCRequest():
                await 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():  # pragma: no branch
                # Exhaustive over JSONRPCMessage, so the no-match arc is unreachable.
                self._resolve_pending(msg.id, msg.error)

    async def _dispatch_request(
        self,
        req: JSONRPCRequest,
        metadata: MessageMetadata,
        on_request: OnRequest,
        sender_ctx: contextvars.Context | None,
    ) -> None:
        progress_token = progress_token_from_params(req.params)
        try:
            transport_ctx = self._transport_builder(metadata)
        except Exception:
            # A raising builder must cost only this message, not the connection.
            logger.exception("transport_builder raised; rejecting request %r", req.id)
            self._spawn(
                self._write_error,
                req.id,
                ErrorData(code=INTERNAL_ERROR, message="transport context unavailable"),
                sender_ctx=sender_ctx,
            )
            return
        dctx = _JSONRPCDispatchContext(
            transport=transport_ctx,
            _dispatcher=self,
            _request_id=req.id,
            message_metadata=metadata,
            _progress_token=progress_token,
        )
        scope = anyio.CancelScope()
        self._corr.enter_inbound(req.id, scope, dctx)
        if req.method in self._inline_methods:
            # Spawn so `sender_ctx` applies, but park the read loop until the
            # handler returns - that's the inline ordering guarantee.
            done = anyio.Event()

            async def _run_inline() -> None:
                try:
                    await self._handle_request(req, dctx, scope, on_request)
                finally:
                    done.set()

            self._spawn(_run_inline, sender_ctx=sender_ctx)
            await done.wait()
        else:
            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:
        """Route one inbound notification.

        `notifications/cancelled` and `notifications/progress` are intercepted
        here (they correlate against the correlator's in-flight/pending
        tables) and still teed to `on_notify` afterwards. The caller's
        `on_notify_intercept` then runs in receive order; only unconsumed
        notifications reach the spawned `on_notify`.
        """
        if msg.method == "notifications/cancelled":
            self._corr.peer_cancel(
                cancelled_request_id_from_params(msg.params),
                interrupt=self._peer_cancel_mode == "interrupt",
            )
        elif msg.method == "notifications/progress":
            delivery = self._corr.progress_callback(msg.params)
            if delivery is not None:
                fn, progress, total, message = delivery
                self._spawn(fn, progress, total, message, sender_ctx=sender_ctx)
        if run_notify_intercept(self._on_notify_intercept, msg.method, msg.params):
            return
        try:
            transport_ctx = self._transport_builder(metadata)
        except Exception:
            # Same containment as `_dispatch_request`: drop the notification, keep the loop.
            logger.exception("transport_builder raised; dropping notification %r", msg.method)
            return
        dctx = _JSONRPCDispatchContext(
            transport=transport_ctx, _dispatcher=self, _request_id=None, message_metadata=metadata
        )
        self._spawn(_contained_notify(on_notify), dctx, msg.method, msg.params, sender_ctx=sender_ctx)

    def _resolve_pending(self, request_id: RequestId | None, outcome: Outcome) -> None:
        self._corr.resolve(request_id, outcome)

    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 task that wrote the
        message; `Context.run` makes the spawned handler inherit that context.
        """
        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_raw_request` waiter with `CONNECTION_CLOSED`. Idempotent."""
        self._corr.fan_out_closed()

    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.

        The exception-to-wire policy lives in `RequestCorrelator.serve_inbound`;
        this only binds the wire writes for a stream-pair transport.
        """
        await self._corr.serve_inbound(
            req.id,
            dctx,
            scope,
            partial(on_request, dctx, req.method, req.params),
            write_result=partial(self._write_result, req.id),
            write_error=partial(self._write_error, req.id),
            raise_handler_exceptions=self._raise_handler_exceptions,
        )

    async def _write(self, message: JSONRPCMessage, metadata: MessageMetadata = None) -> None:
        await self._write_stream.send(SessionMessage(message=message, metadata=metadata))

    async def _write_result(self, request_id: RequestId, result: dict[str, Any]) -> None:
        try:
            await self._write(JSONRPCResponse(jsonrpc="2.0", id=request_id, result=result))
        except (anyio.BrokenResourceError, anyio.ClosedResourceError):
            logger.debug("dropped result for %r: write stream closed", request_id)

    async def _write_error(self, request_id: RequestId, error: ErrorData) -> None:
        try:
            await self._write(JSONRPCError(jsonrpc="2.0", id=request_id, error=error))
        except (anyio.BrokenResourceError, anyio.ClosedResourceError):
            logger.debug("dropped error for %r: write stream closed", request_id)

    async def _cancel_outbound(self, request_id: RequestId, reason: str, related_request_id: RequestId | None) -> None:
        # Thread `related_request_id` so streamable HTTP routes the cancel onto
        # the request's own SSE stream instead of a possibly-absent GET stream.
        # `notify` swallows connection-state errors itself, so no guard here.
        await self.notify(
            "notifications/cancelled",
            {"requestId": request_id, "reason": reason},
            _related_request_id=related_request_id,
        )