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,
        on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None,
        on_read_eof: Callable[[], 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`; runs after admission.
            raise_handler_exceptions: Re-raise handler exceptions out of
                `run()` after the error response is written.
            on_stream_exception: Observer for `Exception` items on the read
                stream, awaited inline in the read loop; without it they are
                debug-logged and dropped.
            on_read_eof: Awaited once when the read stream ends, before
                in-flight handlers are cancelled.
        """
        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.on_stream_exception = on_stream_exception
        """Observer for `Exception` items on the read stream; mutable, consulted only inside `run()`."""
        self._on_read_eof = on_read_eof or _no_drain

        self._next_id = 0
        self._pending: dict[RequestId, _Pending] = {}
        self._in_flight: dict[RequestId, _InFlight[TransportT]] = {}
        self._request_tasks = 0
        """Admitted requests not yet finished answering (their answer writes included)."""
        self._all_answered: anyio.Event | None = None
        """Armed while `wait_for_in_flight` waits; set when `_request_tasks` reaches zero."""
        self._on_notify_intercept: OnNotifyIntercept | None = None
        self._admit: Admit | None = None
        self._tg: anyio.abc.TaskGroup | None = None
        self._running = False
        self._closed = 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` (set only by `_JSONRPCDispatchContext`) routes
        mid-handler requests 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.
        if self._closed:
            raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
        if not self._running:
            raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run()")
        opts = opts or {}
        supplied_id = opts.get("request_id")
        if supplied_id is not None:
            request_id: RequestId = supplied_id
            # Same coercion `_resolve_pending` applies to inbound ids; the wire id stays verbatim.
            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")
        else:
            # Mint past any key a supplied id occupies.
            request_id = self._allocate_id()
            while request_id in self._pending:
                request_id = self._allocate_id()
            pending_key = request_id
        out_params = dict(params) if params is not None else {}
        out_meta = dict(out_params.get("_meta") or {})
        on_progress = opts.get("on_progress")
        if on_progress is not None:
            # The request id doubles as the progress token, so `_pending[token]` finds `on_progress` directly.
            out_meta["progressToken"] = request_id
        out_params["_meta"] = out_meta

        # buffer=1: a close signal can arrive before the waiter parks in receive().
        send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
        pending = _Pending(send=send, receive=receive, on_progress=on_progress)
        self._pending[pending_key] = pending

        plan = _plan_outbound(_related_request_id, opts)
        # A write interrupted by cancellation may still have delivered, so a started
        # write counts as issued (spec MUST: only issued requests may be cancelled).
        request_write_started = False
        timeout_armed = False

        target = out_params.get("name")
        span_name = f"MCP send {method}{f' {target}' if isinstance(target, str) else ''}"
        # TODO(maxisbey): move the otel span + inject into an outbound
        # middleware once that seam exists; the dispatcher should not own otel.
        try:
            with otel_span(
                span_name,
                kind=SpanKind.CLIENT,
                attributes={"mcp.method.name": method, "jsonrpc.request.id": str(request_id)},
            ):
                # SEP-414: inject W3C trace context; `_meta` stays on the wire even with a no-op tracer.
                inject_trace_context(out_meta)
                msg = JSONRPCRequest(jsonrpc="2.0", id=request_id, method=method, params=out_params)
                # Surface a pre-existing cancellation while the request provably never started.
                await anyio.lowlevel.checkpoint_if_cancelled()
                request_write_started = True
                try:
                    await self._write(msg, plan.metadata)
                except (anyio.BrokenResourceError, anyio.ClosedResourceError):
                    # Transport tore down before run() noticed EOF; surface the documented contract.
                    raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") from None
                with anyio.fail_after(opts.get("timeout")):
                    timeout_armed = True
                    outcome = await receive.receive()
        except TimeoutError:
            if not timeout_armed:
                # `fail_after` arms only after the write, so this is the transport's
                # own bounded send() failing, not `opts["timeout"]`: propagate it raw.
                raise
            # Courtesy cancel so the peer stops work; unshielded so an outer cancel can interrupt it.
            if plan.cancel_on_abandon:
                await self._final_write(
                    partial(
                        self._cancel_outbound,
                        request_id,
                        f"timed out after {opts.get('timeout')}s",
                        _related_request_id,
                    ),
                    shield=False,
                    timeout=_ABANDON_WRITE_TIMEOUT,
                    describe=f"courtesy cancel for timed-out request {request_id!r}",
                )
            raise MCPError(code=REQUEST_TIMEOUT, message=f"Request {method!r} timed out") from None
        except anyio.get_cancelled_exc_class():
            # Caller cancelled: the shielded helper sends the courtesy cancel before we propagate.
            if plan.cancel_on_abandon and request_write_started:
                await self._final_write(
                    partial(self._cancel_outbound, request_id, "caller cancelled", _related_request_id),
                    shield=True,
                    timeout=_ABANDON_WRITE_TIMEOUT,
                    describe=f"courtesy cancel for caller-cancelled request {request_id!r}",
                )
            raise
        finally:
            # Remove the waiter on every path so a late response is dropped, not leaked.
            self._pending.pop(pending_key, None)
            send.close()
            receive.close()

        if isinstance(outcome, ErrorData):
            raise MCPError(code=outcome.code, message=outcome.message, data=outcome.data)
        return outcome

    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; a closed dispatcher or torn-down
        transport drops it with a debug log instead of raising."""
        if self._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 wait_for_in_flight(self) -> None:
        """Return once every admitted request has finished answering, answer writes included.

        Unbounded; the caller supplies the timeout. Concurrent waiters all wake.
        """
        while self._request_tasks:
            if self._all_answered is None:
                self._all_answered = anyio.Event()
            await self._all_answered.wait()

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

        `admit`, when given, is consulted synchronously in receive order for each
        inbound request before its transport context is built; without it every
        request goes to `on_request`, unheld. `task_status.started()` fires once
        `send_raw_request` is usable. Single-shot: once the loop ends the
        dispatcher stays closed.
        """
        self._on_notify_intercept = on_notify_intercept
        self._admit = admit
        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: the peer is gone. After the owner's drain, wake
                        # blocked senders and revoke open channels.
                        self._running = False
                        self._closed = True
                        await self._on_read_eof()
                        self._revoke_in_flight()
                        self._fan_out_closed()
                    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 fan-out; idempotent.
            self._running = False
            self._closed = True
            self._tg = None
            self._fan_out_closed()
            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 held requests and the `on_stream_exception`
        observer are awaited, so nothing else head-of-line blocks 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)
        # Admission runs first, in the read loop, in receive order; a raising admission
        # or transport builder costs only this message, never the connection.
        try:
            admission = self._admit(req.method, req.params) if self._admit is not None else Admission(on_request)
        except Exception:
            logger.exception("admission raised; rejecting request %r", req.id)
            self._reject(req.id, ErrorData(code=INTERNAL_ERROR, message="Internal server error"), sender_ctx)
            return
        try:
            transport_ctx = self._transport_builder(metadata)
        except Exception:
            logger.exception("transport_builder raised; rejecting request %r", req.id)
            self._reject(req.id, ErrorData(code=INTERNAL_ERROR, message="transport context unavailable"), sender_ctx)
            return
        dctx = _JSONRPCDispatchContext(
            transport=transport_ctx,
            _target=self._wire(transport_ctx, req.id),
            _request_id=req.id,
            _cancellable=self._in_flight,
            message_metadata=metadata,
            _progress_token=progress_token,
        )
        scope = anyio.CancelScope()
        # The admitted handler is invoked here; the awaitable it returns is the body.
        try:
            body = admission.handler(dctx, req.method, req.params)
        except Exception:
            logger.exception("handler raised during admission; rejecting request %r", req.id)
            self._spawn(
                dctx.reply_error, ErrorData(code=INTERNAL_ERROR, message="Internal server error"), sender_ctx=sender_ctx
            )
            return
        # TODO(maxisbey): duplicate ids blind-overwrite (v1/TS parity); revisit
        # rejecting with INVALID_REQUEST. Key coerced so a stringified
        # `notifications/cancelled` id still correlates.
        self._in_flight[coerce_request_id(req.id)] = _InFlight(scope=scope, dctx=dctx)
        self._request_tasks += 1
        if admission.hold:
            # Spawn so `sender_ctx` applies, but park the read loop until the
            # request finishes answering.
            done = anyio.Event()

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

            self._spawn(_run_held, sender_ctx=sender_ctx)
            await done.wait()
        else:
            self._spawn(self._handle_request, req, dctx, scope, body, 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 correlated here
        first; then the caller's `on_notify_intercept` runs in receive order, and
        only unconsumed notifications are admitted through `on_notify`.
        """
        if msg.method == "notifications/cancelled":
            rid = cancelled_request_id_from_params(msg.params)
            if rid is not None and (in_flight := self._in_flight.get(coerce_request_id(rid))) is not None:
                in_flight.dctx.cancel_requested.set()
                if self._peer_cancel_mode == "interrupt":
                    # The peer withdrew the request: revoke its channel, then interrupt the work.
                    in_flight.dctx.close()
                    in_flight.scope.cancel()
        elif msg.method == "notifications/progress":
            match msg.params:
                case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if (
                    not isinstance(token, bool)
                    and not isinstance(progress, bool)
                    and (pending := self._pending.get(coerce_request_id(token))) is not None
                    and pending.on_progress is not None
                ):
                    total = msg.params.get("total")
                    message = msg.params.get("message")
                    self._spawn(
                        _shielded_progress(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
        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,
            _target=self._wire(transport_ctx, None),
            _request_id=None,
            _cancellable=self._in_flight,
            message_metadata=metadata,
        )
        # Admission mirrors requests (see `OnNotify`).
        try:
            body = on_notify(dctx, msg.method, msg.params)
        except Exception:
            logger.exception("on_notify raised during admission; dropping notification %r", msg.method)
            return
        self._spawn(self._run_notification, msg.method, body, sender_ctx=sender_ctx)

    def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None:
        pending = self._pending.get(coerce_request_id(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 under the sender's contextvars
        context (ASGI middleware sets contextvars on the task that wrote the message)."""
        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`; synchronous, 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,
        body: Awaitable[dict[str, Any]],
    ) -> None:
        """Run one admitted request's `body` inside its cancel scope and answer through `dctx`.

        The exception-to-wire boundary for the body: handler exceptions become
        the request's error frame here. Holds `_request_tasks` until the answer
        write is done, which is what a drain waits for.
        """
        try:
            with scope:
                result = await body
                await dctx.reply(result)
        except anyio.get_cancelled_exc_class():
            # Outer cancel (shutdown): answer a peer that may still be waiting; a
            # spent or revoked channel drops it. Shielded: bare awaits re-raise here.
            await self._final_write(
                partial(dctx.reply_error, ErrorData(code=CONNECTION_CLOSED, message="Connection closed")),
                shield=True,
                timeout=_SHUTDOWN_WRITE_TIMEOUT,
                describe=f"shutdown error response for request {req.id!r}",
            )
            raise
        except Exception as e:
            error = handler_exception_to_error_data(e)
            if error is not None:
                await dctx.reply_error(error)
            else:
                logger.exception("handler for %r raised", req.method)
                # Only the generic error reaches the wire; the exception text stays in the log.
                await dctx.reply_error(ErrorData(code=INTERNAL_ERROR, message="Internal server error"))
                if self._raise_handler_exceptions:
                    raise
        finally:
            self._release_request_task()

    async def _run_notification(self, method: str, body: Awaitable[None]) -> None:
        """Run one admitted notification's `body`, containing a raise to that message."""
        try:
            await body
        except Exception:
            logger.exception("notification handler for %r raised", method)

    def _release_request_task(self) -> None:
        """A request task has finished: wake `wait_for_in_flight` when it was the last one."""
        self._request_tasks -= 1
        if not self._request_tasks and (answered := self._all_answered) is not None:
            self._all_answered = None
            answered.set()

    def _revoke_in_flight(self) -> None:
        """The peer is gone: revoke every open request's channel so nothing further is written."""
        for entry in self._in_flight.values():
            entry.dctx.close()

    def _allocate_id(self) -> int:
        self._next_id += 1
        return self._next_id

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

    async def _write_frame(self, frame: JSONRPCMessage) -> None:
        """Write one pre-built frame onto the connection; a torn-down transport drops it."""
        try:
            await self._write(frame)
        except (anyio.BrokenResourceError, anyio.ClosedResourceError):
            logger.debug("dropped %s: write stream closed", type(frame).__name__)

    def _reject(self, request_id: RequestId, error: ErrorData, sender_ctx: contextvars.Context | None) -> None:
        """Answer a request the read loop could not admit; the error frame is its whole fate."""
        self._spawn(self._write_frame, JSONRPCError(jsonrpc="2.0", id=request_id, error=error), sender_ctx=sender_ctx)

    def _wire(self, transport: TransportT, rid: RequestId | None) -> _Wire[TransportT]:
        """The live write target of a message arriving on this connection."""
        return _Wire(self, self._write_frame, transport, rid)

    async def _final_write(
        self,
        write: Callable[[], Awaitable[None]],
        *,
        shield: bool,
        timeout: float,
        describe: str,
    ) -> None:
        """One last write under the shared abandon/teardown policy: bounded so a wedged
        transport cannot hang; `shield=True` for arms already inside a cancelled scope."""
        with anyio.move_on_after(timeout, shield=shield) as scope:
            await write()
        if scope.cancelled_caught:
            logger.warning("%s gave up: transport write blocked", describe)

    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; `notify` swallows connection-state errors itself.
        await self.notify(
            "notifications/cancelled",
            {"requestId": request_id, "reason": reason},
            _related_request_id=related_request_id,
        )