742 lines
33 KiB
Python
742 lines
33 KiB
Python
"""JSON-RPC `Dispatcher` implementation.
|
|
|
|
Consumes the existing `SessionMessage`-based stream contract that all current
|
|
transports (stdio, SSE, streamable HTTP) speak. Owns request-id correlation,
|
|
the receive loop, per-request task isolation, cancellation/progress wiring, and
|
|
the single exception-to-wire boundary.
|
|
|
|
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 Awaitable, Callable, Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Generic, Literal, TypeVar, cast, overload
|
|
|
|
import anyio
|
|
import anyio.abc
|
|
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
|
from opentelemetry.trace import SpanKind
|
|
from pydantic import ValidationError
|
|
|
|
from mcp.shared._otel import inject_trace_context, otel_span
|
|
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
|
from mcp.shared.dispatcher import CallOptions, Dispatcher, OnNotify, OnRequest, ProgressFnT
|
|
from mcp.shared.exceptions import MCPError, NoBackChannelError
|
|
from mcp.shared.message import (
|
|
ClientMessageMetadata,
|
|
MessageMetadata,
|
|
ServerMessageMetadata,
|
|
SessionMessage,
|
|
)
|
|
from mcp.shared.transport_context import TransportContext
|
|
from mcp.types import (
|
|
CONNECTION_CLOSED,
|
|
INTERNAL_ERROR,
|
|
INVALID_PARAMS,
|
|
REQUEST_CANCELLED,
|
|
REQUEST_TIMEOUT,
|
|
ErrorData,
|
|
JSONRPCError,
|
|
JSONRPCMessage,
|
|
JSONRPCNotification,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
ProgressToken,
|
|
RequestId,
|
|
)
|
|
|
|
__all__ = ["JSONRPCDispatcher"]
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TransportT = TypeVar("TransportT", bound=TransportContext)
|
|
|
|
PeerCancelMode = Literal["interrupt", "signal"]
|
|
"""How inbound `notifications/cancelled` is applied to a running handler.
|
|
|
|
`"interrupt"` (default) cancels the handler's scope. `"signal"` only sets
|
|
`ctx.cancel_requested` and lets the handler observe it cooperatively.
|
|
"""
|
|
|
|
|
|
def _coerce_id(request_id: RequestId) -> RequestId:
|
|
"""Coerce a string request ID to int when it's a valid int literal.
|
|
|
|
`_allocate_id` only ever produces `int` keys for `_pending`, but a peer
|
|
may echo the ID back as a JSON string. The TypeScript SDK and `BaseSession`
|
|
both perform this coercion at lookup time so the response still correlates.
|
|
"""
|
|
if isinstance(request_id, str):
|
|
try:
|
|
return int(request_id)
|
|
except ValueError:
|
|
pass
|
|
return request_id
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _Pending:
|
|
"""An outbound request awaiting its response."""
|
|
|
|
send: MemoryObjectSendStream[dict[str, Any] | ErrorData]
|
|
receive: MemoryObjectReceiveStream[dict[str, Any] | ErrorData]
|
|
on_progress: ProgressFnT | None = None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _InFlight(Generic[TransportT]):
|
|
"""An inbound request currently being handled."""
|
|
|
|
scope: anyio.CancelScope
|
|
dctx: _JSONRPCDispatchContext[TransportT]
|
|
|
|
|
|
@dataclass
|
|
class _JSONRPCDispatchContext(Generic[TransportT]):
|
|
"""Concrete `DispatchContext` produced for each inbound JSON-RPC message."""
|
|
|
|
transport: TransportT
|
|
_dispatcher: JSONRPCDispatcher[TransportT]
|
|
_request_id: RequestId | None
|
|
message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework
|
|
"""The transport-attached `SessionMessage.metadata` for this inbound message.
|
|
|
|
Carries `ServerMessageMetadata` (HTTP request, SSE stream-close callbacks)
|
|
that the server lifts onto its request context. `None` for transports
|
|
that attach nothing.
|
|
"""
|
|
_progress_token: ProgressToken | None = None
|
|
_closed: bool = False
|
|
cancel_requested: anyio.Event = field(default_factory=anyio.Event)
|
|
|
|
@property
|
|
def request_id(self) -> RequestId | None:
|
|
return self._request_id
|
|
|
|
@property
|
|
def can_send_request(self) -> bool:
|
|
return self.transport.can_send_request and not self._closed
|
|
|
|
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
|
if self._closed:
|
|
logger.debug("dropped %s: dispatch context closed", method)
|
|
return
|
|
await self._dispatcher.notify(method, params, _related_request_id=self._request_id)
|
|
|
|
async def send_raw_request(
|
|
self,
|
|
method: str,
|
|
params: Mapping[str, Any] | None,
|
|
opts: CallOptions | None = None,
|
|
) -> dict[str, Any]:
|
|
if not self.can_send_request:
|
|
raise NoBackChannelError(method)
|
|
return await self._dispatcher.send_raw_request(method, params, opts, _related_request_id=self._request_id)
|
|
|
|
async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
|
if self._progress_token is None:
|
|
return
|
|
params: dict[str, Any] = {"progressToken": self._progress_token, "progress": progress}
|
|
if total is not None:
|
|
params["total"] = total
|
|
if message is not None:
|
|
params["message"] = message
|
|
await self.notify("notifications/progress", params)
|
|
|
|
def close(self) -> None:
|
|
self._closed = True
|
|
|
|
|
|
def _default_transport_builder(_meta: MessageMetadata) -> TransportContext:
|
|
return TransportContext(kind="jsonrpc", can_send_request=True)
|
|
|
|
|
|
def _shielded_progress(fn: ProgressFnT) -> ProgressFnT:
|
|
"""Wrap a user progress callback so it can't crash the dispatcher.
|
|
|
|
The callback runs as a bare task in the dispatcher's task group; an
|
|
uncaught exception would cancel every sibling (the read loop and all
|
|
in-flight requests). Swallow and log instead, matching the previous
|
|
receive-loop's behavior.
|
|
"""
|
|
|
|
async def _wrapped(progress: float, total: float | None, message: str | None) -> None:
|
|
try:
|
|
await fn(progress, total, message)
|
|
except Exception:
|
|
logger.exception("progress callback raised")
|
|
|
|
return _wrapped
|
|
|
|
|
|
def _outbound_metadata(related_request_id: RequestId | None, opts: CallOptions | None) -> MessageMetadata:
|
|
"""Choose the `SessionMessage.metadata` for an outgoing request/notification.
|
|
|
|
`ServerMessageMetadata` tags a server-to-client message with the inbound
|
|
request it belongs to (so streamable-HTTP can route it onto that request's
|
|
SSE stream). `ClientMessageMetadata` carries resumption hints to the
|
|
client transport. `None` is the common case.
|
|
|
|
`SessionMessage.metadata` carries exactly one of these, so when
|
|
`related_request_id` is set it takes precedence and any resumption hints
|
|
in `opts` are dropped (with a debug log): requests made from a dispatch
|
|
context are routed onto the inbound request's stream, not resumed.
|
|
"""
|
|
if related_request_id is not None:
|
|
if opts and (opts.get("resumption_token") is not None or opts.get("on_resumption_token") is not None):
|
|
logger.debug(
|
|
"dropping resumption hints: related_request_id %r takes precedence on metadata", related_request_id
|
|
)
|
|
return ServerMessageMetadata(related_request_id=related_request_id)
|
|
if opts:
|
|
token = opts.get("resumption_token")
|
|
on_token = opts.get("on_resumption_token")
|
|
if token is not None or on_token is not None:
|
|
return ClientMessageMetadata(resumption_token=token, on_resumption_token_update=on_token)
|
|
return None
|
|
|
|
|
|
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__(
|
|
self: JSONRPCDispatcher[TransportContext],
|
|
read_stream: ReadStream[SessionMessage | Exception],
|
|
write_stream: WriteStream[SessionMessage],
|
|
*,
|
|
peer_cancel_mode: PeerCancelMode = "interrupt",
|
|
raise_handler_exceptions: bool = False,
|
|
inline_methods: frozenset[str] = frozenset(),
|
|
) -> None: ...
|
|
@overload
|
|
def __init__(
|
|
self,
|
|
read_stream: ReadStream[SessionMessage | Exception],
|
|
write_stream: WriteStream[SessionMessage],
|
|
*,
|
|
transport_builder: Callable[[MessageMetadata], TransportT],
|
|
peer_cancel_mode: PeerCancelMode = "interrupt",
|
|
raise_handler_exceptions: bool = False,
|
|
inline_methods: frozenset[str] = frozenset(),
|
|
) -> None: ...
|
|
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(),
|
|
) -> None:
|
|
self._read_stream = read_stream
|
|
self._write_stream = write_stream
|
|
# 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[[MessageMetadata], TransportT]",
|
|
transport_builder or _default_transport_builder,
|
|
)
|
|
self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode
|
|
self._raise_handler_exceptions = raise_handler_exceptions
|
|
# Request methods handled inline in the read loop (awaited before the
|
|
# next message is dequeued) instead of spawned concurrently. Use for
|
|
# methods whose side effects must be observable to the next message,
|
|
# e.g. `initialize`, so a pipelined follow-up sees the initialized state.
|
|
# Only suitable for handlers that complete quickly, since inline handling
|
|
# blocks dequeuing; a handler that awaits the peer (`send_raw_request`)
|
|
# while inline will deadlock because the parked read loop cannot dequeue
|
|
# the response.
|
|
self._inline_methods = inline_methods
|
|
|
|
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_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` when a
|
|
handler makes a server-to-client request mid-flight; it routes the
|
|
outgoing message onto the correct per-request SSE stream (SHTTP) via
|
|
`ServerMessageMetadata`. Top-level callers leave it `None`.
|
|
|
|
Raises:
|
|
MCPError: The peer responded with a JSON-RPC error; or
|
|
`REQUEST_TIMEOUT` if `opts["timeout"]` elapsed; or
|
|
`CONNECTION_CLOSED` if the dispatcher shut down while
|
|
awaiting the response.
|
|
RuntimeError: Called before `run()` has started or after it has
|
|
finished.
|
|
"""
|
|
if not self._running:
|
|
raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run() / after close")
|
|
opts = opts or {}
|
|
request_id = self._allocate_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 caller wants progress updates. The spec mechanism is: include
|
|
# `_meta.progressToken` on the request; the peer echoes that token on
|
|
# any `notifications/progress` it sends. We use the request id as the
|
|
# token so the receive loop can find this `_Pending.on_progress` by
|
|
# `_pending[token]` without a second lookup table.
|
|
out_meta["progressToken"] = request_id
|
|
out_params["_meta"] = out_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
|
|
|
|
metadata = _outbound_metadata(_related_request_id, opts)
|
|
target = out_params.get("name")
|
|
span_name = f"MCP send {method}{f' {target}' if isinstance(target, str) else ''}"
|
|
# TODO(maxisbey): the otel span + inject below mirror
|
|
# BaseSession.send_request for parity. They belong in an outbound
|
|
# middleware (symmetric with otel_middleware on the inbound side) 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)},
|
|
):
|
|
# Inject W3C trace context into _meta (SEP-414). With a no-op
|
|
# tracer this writes nothing, but `_meta` itself is still
|
|
# present on the wire (and the interaction suite pins that).
|
|
inject_trace_context(out_meta)
|
|
msg = JSONRPCRequest(jsonrpc="2.0", id=request_id, method=method, params=out_params)
|
|
await self._write(msg, metadata)
|
|
with anyio.fail_after(opts.get("timeout")):
|
|
outcome = await receive.receive()
|
|
except TimeoutError:
|
|
# Spec-recommended courtesy: tell the peer we've given up so it can
|
|
# stop work and free resources. v1's BaseSession.send_request does
|
|
# NOT do this; it's new behaviour.
|
|
await self._cancel_outbound(request_id, f"timed out after {opts.get('timeout')}s", _related_request_id)
|
|
raise MCPError(code=REQUEST_TIMEOUT, message=f"Request {method!r} timed out") from None
|
|
except anyio.get_cancelled_exc_class():
|
|
# Our caller's scope was cancelled. We're already inside a cancelled
|
|
# scope, so any bare `await` here re-raises immediately - shield to
|
|
# let the courtesy cancel notification go out before we propagate.
|
|
with anyio.CancelScope(shield=True):
|
|
await self._cancel_outbound(request_id, "caller cancelled", _related_request_id)
|
|
raise
|
|
finally:
|
|
# Always remove the waiter, even on cancel/timeout, so a late
|
|
# response from the peer (race) hits a closed stream and is dropped
|
|
# in `_dispatch` rather than leaking.
|
|
self._pending.pop(request_id, 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,
|
|
*,
|
|
_related_request_id: RequestId | None = None,
|
|
) -> None:
|
|
# Leave `params` unset (not explicitly None) when there are none:
|
|
# transports serialize with `exclude_unset=True`, and an explicit None
|
|
# would survive as `"params": null`, which JSON-RPC 2.0 forbids and
|
|
# strict peers (e.g. the TypeScript SDK's zod schemas) reject.
|
|
if params is not None:
|
|
msg = JSONRPCNotification(jsonrpc="2.0", method=method, params=dict(params))
|
|
else:
|
|
msg = JSONRPCNotification(jsonrpc="2.0", method=method)
|
|
await self._write(msg, _outbound_metadata(_related_request_id, 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 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_raw_request`
|
|
is usable.
|
|
"""
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
self._tg = tg
|
|
self._running = True
|
|
task_status.started()
|
|
try:
|
|
async with self._read_stream, self._write_stream:
|
|
try:
|
|
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
|
|
)
|
|
await self._dispatch(item, on_request, on_notify, sender_ctx)
|
|
except anyio.ClosedResourceError:
|
|
# The transport closed our receive end and we looped
|
|
# back to `__anext__` on the now-closed stream
|
|
# (stateless SHTTP teardown). Same as EOF.
|
|
logger.debug("read stream closed by transport; treating as EOF")
|
|
# Read stream EOF: wake any blocked `send_raw_request` waiters
|
|
# (callers outside this task group) with CONNECTION_CLOSED.
|
|
self._running = False
|
|
self._fan_out_closed()
|
|
finally:
|
|
# Transport closed: cancel in-flight handlers. Without this
|
|
# the task-group join waits for them, and a handler that
|
|
# outlives its caller (its request timed out client-side, or
|
|
# the client disconnected mid-call) would keep `run()` from
|
|
# returning forever. Same behaviour as `Server.run()` before
|
|
# the dispatcher rework.
|
|
tg.cancel_scope.cancel()
|
|
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()
|
|
|
|
async def _dispatch(
|
|
self,
|
|
item: SessionMessage | Exception,
|
|
on_request: OnRequest,
|
|
on_notify: OnNotify,
|
|
sender_ctx: contextvars.Context | None,
|
|
) -> None:
|
|
"""Route one inbound item.
|
|
|
|
Everything here is `send_nowait` or `_spawn`; the only `await` is for
|
|
`inline_methods` requests, which deliberately block dequeuing until
|
|
handled. Any other `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():
|
|
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
|
|
# `id` may be None per JSON-RPC (parse error before id known).
|
|
# The match is exhaustive over JSONRPCMessage; the no-match arc
|
|
# on this final case 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: ProgressToken | None
|
|
match req.params:
|
|
# The bool guard matters: `int()` patterns match bool (a subclass),
|
|
# and `True == 1` would alias dict lookups to request id 1.
|
|
case {"_meta": {"progressToken": str() | int() as progress_token}} if not isinstance(progress_token, bool):
|
|
pass
|
|
case _:
|
|
progress_token = None
|
|
try:
|
|
transport_ctx = self._transport_builder(metadata)
|
|
except Exception:
|
|
# Containment boundary for the user-supplied builder: a raising
|
|
# builder must cost only this message, not the whole connection
|
|
# (the exception would otherwise escape into run()'s read loop).
|
|
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()
|
|
# TODO(maxisbey): the spec puts request-id uniqueness on the sender;
|
|
# neither v1 nor the TS SDK guards a duplicate id here, so for now we
|
|
# blind-overwrite (parity). Revisit rejecting with INVALID_REQUEST.
|
|
# Coerced key so `notifications/cancelled` correlates regardless of
|
|
# whether the peer stringifies the id between request and cancel
|
|
# (`_dispatch_notification` coerces at lookup; responses still echo
|
|
# `req.id` verbatim).
|
|
self._in_flight[_coerce_id(req.id)] = _InFlight(scope=scope, dctx=dctx)
|
|
if req.method in self._inline_methods:
|
|
# Spawn (so `sender_ctx` applies, matching the concurrent path) but
|
|
# park the read loop until the handler returns; that's the inline
|
|
# ordering guarantee. Because the read loop is parked, a handler
|
|
# that awaits the peer here (e.g. `dctx.send_raw_request`) will
|
|
# deadlock: the response can never be dequeued.
|
|
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 because they correlate against JSON-RPC request IDs - the
|
|
`_in_flight` / `_pending` tables this layer owns - so no higher layer
|
|
can act on them. Both are still teed to `on_notify` afterwards, so
|
|
middleware and registered notification handlers observe every inbound
|
|
notification. See the module docstring for the design rationale.
|
|
"""
|
|
if msg.method == "notifications/cancelled":
|
|
match msg.params:
|
|
# The bool guards here and below matter: `int()` patterns match
|
|
# bool (a subclass), and `True == 1` would alias the dict lookup
|
|
# to the entry keyed by request id 1.
|
|
case {"requestId": str() | int() as rid} if (
|
|
not isinstance(rid, bool) and (in_flight := self._in_flight.get(_coerce_id(rid))) is not None
|
|
):
|
|
in_flight.dctx.cancel_requested.set()
|
|
if self._peer_cancel_mode == "interrupt":
|
|
in_flight.scope.cancel()
|
|
case _:
|
|
pass
|
|
# fall through: cancelled is also teed to on_notify so middleware
|
|
# and registered handlers can observe it (matches DirectDispatcher,
|
|
# which forwards every notification).
|
|
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_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
|
|
# fall through: progress is also teed to on_notify
|
|
try:
|
|
transport_ctx = self._transport_builder(metadata)
|
|
except Exception:
|
|
# Same containment boundary as `_dispatch_request`: a raising
|
|
# builder drops this notification instead of killing the read 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(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(_coerce_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, 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_raw_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.
|
|
|
|
This is the single exception-to-wire boundary: handler exceptions are
|
|
caught here and serialized to `JSONRPCError`. Nothing above this in
|
|
the stack constructs wire errors.
|
|
"""
|
|
try:
|
|
with scope:
|
|
try:
|
|
result = await on_request(dctx, req.method, req.params)
|
|
finally:
|
|
# Handler done: close the back-channel (detached work that
|
|
# later calls `dctx.send_raw_request()` should see
|
|
# `NoBackChannelError`) and drop from `_in_flight` so a
|
|
# late `notifications/cancelled` is a no-op rather than
|
|
# racing the result write below. No checkpoint between
|
|
# handler return and the pop, so the cancel can't
|
|
# interleave there.
|
|
dctx.close()
|
|
self._in_flight.pop(_coerce_id(req.id), None)
|
|
await self._write_result(req.id, result)
|
|
if scope.cancel_called:
|
|
# Peer-cancel: `_dispatch_notification` cancelled this scope
|
|
# while the handler was running. anyio swallows a scope's *own*
|
|
# cancel at __exit__, so execution lands here rather than the
|
|
# `except cancelled` arm below.
|
|
# TODO(maxisbey): spec says SHOULD NOT respond after cancel.
|
|
# The existing server always has, so match that for now.
|
|
await self._write_error(req.id, ErrorData(code=0, message="Request cancelled"))
|
|
except anyio.get_cancelled_exc_class():
|
|
# Outer-cancel: run()'s task group is shutting down. Any bare
|
|
# `await` here re-raises immediately, so shield the courtesy write.
|
|
with anyio.CancelScope(shield=True):
|
|
await self._write_error(req.id, ErrorData(code=REQUEST_CANCELLED, message="Request cancelled"))
|
|
raise
|
|
except MCPError as e:
|
|
await self._write_error(req.id, e.error)
|
|
except ValidationError:
|
|
# TODO(maxisbey): data="" is pinned compat with the existing
|
|
# server (which never leaked pydantic error text onto the wire).
|
|
# Consider putting the validation detail in `data` once the
|
|
# interaction suite's divergence entry is resolved.
|
|
await self._write_error(
|
|
req.id, ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="")
|
|
)
|
|
except Exception as e:
|
|
logger.exception("handler for %r raised", req.method)
|
|
# TODO(maxisbey): code=0 is pinned compat with the existing
|
|
# server's `_handle_request`. JSON-RPC says INTERNAL_ERROR
|
|
# (-32603); revisit once the suite's divergence entry is resolved.
|
|
await self._write_error(req.id, ErrorData(code=0, message=str(e)))
|
|
if self._raise_handler_exceptions:
|
|
raise
|
|
# No outer `_in_flight` pop here: the inner `finally` above already
|
|
# removes the entry on every path out of the handler, and a second
|
|
# pop after the awaited response writes could evict a newer request
|
|
# that reused the id during that window.
|
|
|
|
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_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 same per-request SSE stream as the request it cancels; without it
|
|
# the notification falls through to the standalone GET stream and is
|
|
# dropped when no GET stream is open.
|
|
try:
|
|
await self.notify(
|
|
"notifications/cancelled",
|
|
{"requestId": request_id, "reason": reason},
|
|
_related_request_id=related_request_id,
|
|
)
|
|
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
|
|
pass
|