335 lines
14 KiB
Python
335 lines
14 KiB
Python
"""In-memory `Dispatcher` that wires two peers together with no transport.
|
|
|
|
`DirectDispatcher` is the simplest possible `Dispatcher` implementation: a
|
|
request on one side directly invokes the other side's `on_request`. There is no
|
|
serialization, no JSON-RPC framing, and no streams. It exists to:
|
|
|
|
* prove the `Dispatcher` Protocol is implementable without JSON-RPC
|
|
* provide a fast substrate for testing the layers above the dispatcher
|
|
(`ServerRunner`, `Context`, `Connection`) without wire-level moving parts
|
|
* embed a server in-process when the JSON-RPC overhead is unnecessary
|
|
|
|
Like `JSONRPCDispatcher`, this is an exception-to-error boundary: a handler
|
|
exception surfaces to the caller as `MCPError`. The `raise_handler_exceptions`
|
|
knob controls whether unmapped exceptions are sanitized (matching the wire
|
|
path) or chained as ``__cause__`` for in-process debugging.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections.abc import Awaitable, Callable, Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import anyio
|
|
import anyio.abc
|
|
from mcp_types import CONNECTION_CLOSED, INTERNAL_ERROR, INVALID_PARAMS, REQUEST_TIMEOUT, RequestId
|
|
from pydantic import ValidationError
|
|
|
|
from mcp.shared._compat import resync_tracer
|
|
from mcp.shared.dispatcher import (
|
|
CallOptions,
|
|
OnNotify,
|
|
OnNotifyIntercept,
|
|
OnRequest,
|
|
ProgressFnT,
|
|
coerce_request_id,
|
|
run_notify_intercept,
|
|
)
|
|
from mcp.shared.exceptions import MCPError, NoBackChannelError
|
|
from mcp.shared.message import MessageMetadata
|
|
from mcp.shared.transport_context import TransportContext
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
__all__ = ["DirectDispatcher", "create_direct_dispatcher_pair"]
|
|
|
|
DIRECT_TRANSPORT_KIND = "direct"
|
|
|
|
|
|
_Request = Callable[[str, Mapping[str, Any] | None, CallOptions | None], Awaitable[dict[str, Any]]]
|
|
_Notify = Callable[[str, Mapping[str, Any] | None], Awaitable[None]]
|
|
|
|
|
|
@dataclass
|
|
class _DirectDispatchContext:
|
|
"""`DispatchContext` for an inbound request on a `DirectDispatcher`.
|
|
|
|
The back-channel callables target the *originating* side, so a handler's
|
|
`send_raw_request` reaches the peer that made the inbound request.
|
|
"""
|
|
|
|
transport: TransportContext
|
|
_back_request: _Request
|
|
_back_notify: _Notify
|
|
request_id: RequestId | None = None
|
|
"""The caller-supplied `CallOptions["request_id"]`, else a dispatcher-synthesized
|
|
id for requests; `None` for notifications."""
|
|
message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework
|
|
"""Always `None`: in-memory dispatch attaches no transport metadata."""
|
|
_on_progress: ProgressFnT | None = None
|
|
cancel_requested: anyio.Event = field(default_factory=anyio.Event)
|
|
|
|
@property
|
|
def can_send_request(self) -> bool:
|
|
return self.transport.can_send_request
|
|
|
|
async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None:
|
|
await self._back_notify(method, params)
|
|
|
|
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._back_request(method, params, opts)
|
|
|
|
async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
|
if self._on_progress is not None:
|
|
await self._on_progress(progress, total, message)
|
|
|
|
|
|
class DirectDispatcher:
|
|
"""A `Dispatcher` that calls a peer's handlers directly, in-process.
|
|
|
|
Two instances are wired together with `create_direct_dispatcher_pair`; each
|
|
holds a reference to the other. `send_raw_request` on one awaits the peer's
|
|
`on_request`. `run` parks until `close` is called.
|
|
|
|
Lifecycle mirrors `JSONRPCDispatcher`: `send_raw_request` requires `run()`
|
|
to have started, and once a side has closed - via `close()` or `run()`
|
|
ending - `send_raw_request` raises `MCPError` (`CONNECTION_CLOSED`) and
|
|
inbound requests fail the peer's call the same way instead of invoking the
|
|
handler. Notifications are fire-and-forget in both directions: after close
|
|
they are silently dropped.
|
|
"""
|
|
|
|
def __init__(self, transport_ctx: TransportContext, *, raise_handler_exceptions: bool = True):
|
|
self._transport_ctx = transport_ctx
|
|
self._raise_handler_exceptions = raise_handler_exceptions
|
|
self._peer: DirectDispatcher | None = None
|
|
self._on_request: OnRequest | None = None
|
|
self._on_notify: OnNotify | None = None
|
|
self._on_notify_intercept: OnNotifyIntercept | None = None
|
|
self._next_id = 0
|
|
self._in_flight_ids: set[RequestId] = set()
|
|
self._ready = anyio.Event()
|
|
self._close_event = anyio.Event()
|
|
self._running = False
|
|
self._closed = False
|
|
|
|
def connect_to(self, peer: DirectDispatcher) -> None:
|
|
self._peer = peer
|
|
|
|
async def send_raw_request(
|
|
self,
|
|
method: str,
|
|
params: Mapping[str, Any] | None,
|
|
opts: CallOptions | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Send a request by invoking the peer's `on_request` directly.
|
|
|
|
Raises:
|
|
MCPError: The peer's handler raised; `REQUEST_TIMEOUT` if
|
|
`opts["timeout"]` elapsed; `CONNECTION_CLOSED` if either
|
|
side has closed.
|
|
RuntimeError: Called before `run()`.
|
|
"""
|
|
if self._peer is None:
|
|
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
|
|
# Post-close sends get the same CONNECTION_CLOSED contract as JSONRPCDispatcher.
|
|
if self._closed:
|
|
raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
|
|
if not self._running:
|
|
raise RuntimeError("DirectDispatcher.send_raw_request called before run()")
|
|
return await self._peer._dispatch_request(method, params, opts)
|
|
|
|
async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None:
|
|
"""Send a notification by invoking the peer's `on_notify` directly.
|
|
|
|
Fire-and-forget: usable before `run()` (delivery waits for the peer to
|
|
start), and after close it is silently dropped, matching
|
|
`JSONRPCDispatcher.notify`. `opts` is accepted for `Dispatcher`
|
|
conformance; there is no HTTP layer here so `headers` is ignored.
|
|
"""
|
|
if self._peer is None:
|
|
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
|
|
if self._closed:
|
|
logger.debug("dropped notification %r on closed DirectDispatcher", method)
|
|
return
|
|
await self._peer._dispatch_notify(method, params)
|
|
|
|
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:
|
|
"""Mark this side ready and park until `close()` is called.
|
|
|
|
Single-shot, like `JSONRPCDispatcher.run`: once it returns the
|
|
dispatcher stays closed and cannot be restarted.
|
|
"""
|
|
try:
|
|
self._on_request = on_request
|
|
self._on_notify = on_notify
|
|
self._on_notify_intercept = on_notify_intercept
|
|
self._running = True
|
|
self._ready.set()
|
|
task_status.started()
|
|
await self._close_event.wait()
|
|
finally:
|
|
self._running = False
|
|
self._closed = True
|
|
# run() may end via cancellation without close() ever being
|
|
# called; setting the event wakes `_wait_ready` waiters so they
|
|
# observe the closed state instead of parking forever.
|
|
self._close_event.set()
|
|
|
|
def close(self) -> None:
|
|
self._closed = True
|
|
self._close_event.set()
|
|
|
|
def _make_context(
|
|
self, on_progress: ProgressFnT | None = None, request_id: RequestId | None = None
|
|
) -> _DirectDispatchContext:
|
|
assert self._peer is not None
|
|
peer = self._peer
|
|
return _DirectDispatchContext(
|
|
transport=self._transport_ctx,
|
|
_back_request=lambda m, p, o: peer._dispatch_request(m, p, o),
|
|
_back_notify=lambda m, p: peer._dispatch_notify(m, p),
|
|
request_id=request_id,
|
|
_on_progress=on_progress,
|
|
)
|
|
|
|
async def _wait_ready(self) -> None:
|
|
"""Park until `run()` has started, waking early if this side closes.
|
|
|
|
Raises:
|
|
MCPError: `CONNECTION_CLOSED` if this side has closed.
|
|
"""
|
|
if not self._ready.is_set() and not self._close_event.is_set():
|
|
async with anyio.create_task_group() as tg:
|
|
|
|
async def wake_on(event: anyio.Event) -> None:
|
|
await event.wait()
|
|
tg.cancel_scope.cancel()
|
|
|
|
tg.start_soon(wake_on, self._ready)
|
|
tg.start_soon(wake_on, self._close_event)
|
|
if self._closed:
|
|
raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
|
|
|
|
async def _dispatch_request(
|
|
self,
|
|
method: str,
|
|
params: Mapping[str, Any] | None,
|
|
opts: CallOptions | None,
|
|
) -> dict[str, Any]:
|
|
opts = opts or {}
|
|
try:
|
|
with anyio.fail_after(opts.get("timeout")):
|
|
# Inside the timeout scope, so a configured timeout also bounds
|
|
# waiting on a peer whose run() has not started yet.
|
|
await self._wait_ready()
|
|
assert self._on_request is not None
|
|
supplied_id = opts.get("request_id")
|
|
if supplied_id is not None:
|
|
request_id: RequestId = supplied_id
|
|
# Collisions use the same coerced domain as JSONRPCDispatcher's
|
|
# pending keys, so this in-memory stand-in raises for exactly
|
|
# the ids the wire dispatcher would; the context still sees
|
|
# the verbatim value.
|
|
in_flight_key = coerce_request_id(request_id)
|
|
if in_flight_key in self._in_flight_ids:
|
|
raise ValueError(f"request id {request_id!r} is already in flight")
|
|
else:
|
|
# Synthesize an id (the DispatchContext contract reserves None
|
|
# for notifications), minting past any key a supplied id
|
|
# occupies: the collision error is reserved for the caller
|
|
# who actually chose the id.
|
|
self._next_id += 1
|
|
while self._next_id in self._in_flight_ids:
|
|
self._next_id += 1
|
|
request_id = self._next_id
|
|
in_flight_key = request_id
|
|
self._in_flight_ids.add(in_flight_key)
|
|
dctx = self._make_context(on_progress=opts.get("on_progress"), request_id=request_id)
|
|
try:
|
|
return await self._on_request(dctx, method, params)
|
|
except MCPError:
|
|
raise
|
|
except ValidationError as e:
|
|
# Same shape JSONRPCDispatcher writes, so runner-over-direct
|
|
# tests see what runner-over-JSONRPC would.
|
|
raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="") from e
|
|
except Exception as e:
|
|
# Single owner of the in-proc exception-to-error policy (mirrors
|
|
# JSONRPCDispatcher / `_streamable_http_modern._to_jsonrpc_response`
|
|
# for the wire paths). True chains the original for in-process
|
|
# debugging; False sanitizes to match the wire path's leak guard.
|
|
if self._raise_handler_exceptions:
|
|
raise MCPError(code=INTERNAL_ERROR, message=str(e)) from e
|
|
logger.exception("request handler raised")
|
|
raise MCPError(code=INTERNAL_ERROR, message="Internal server error") from None
|
|
finally:
|
|
self._in_flight_ids.discard(in_flight_key)
|
|
except TimeoutError:
|
|
raise MCPError(
|
|
code=REQUEST_TIMEOUT,
|
|
message=f"Timed out after {opts.get('timeout')}s waiting for {method!r}",
|
|
) from None
|
|
finally:
|
|
await resync_tracer()
|
|
|
|
async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
|
try:
|
|
await self._wait_ready()
|
|
except MCPError:
|
|
# Notifications are fire-and-forget: a notify to a closed peer is
|
|
# dropped, not raised back into the sender's call.
|
|
logger.debug("dropped notification %r to closed DirectDispatcher", method)
|
|
return
|
|
if run_notify_intercept(self._on_notify_intercept, method, params):
|
|
return
|
|
assert self._on_notify is not None
|
|
dctx = self._make_context()
|
|
await self._on_notify(dctx, method, params)
|
|
|
|
|
|
def create_direct_dispatcher_pair(
|
|
*,
|
|
can_send_request: bool = True,
|
|
headers: Mapping[str, str] | None = None,
|
|
raise_handler_exceptions: bool = True,
|
|
) -> tuple[DirectDispatcher, DirectDispatcher]:
|
|
"""Create two `DirectDispatcher` instances wired to each other.
|
|
|
|
Args:
|
|
can_send_request: Sets `TransportContext.can_send_request` on both
|
|
sides. Pass `False` to simulate a transport with no back-channel.
|
|
headers: Sets `TransportContext.headers` on both sides.
|
|
raise_handler_exceptions: When `True` (the default - this is an
|
|
in-process debugging substrate), an unmapped handler exception
|
|
reaches the caller as `MCPError` with the original chained as
|
|
``__cause__``. When `False` it is sanitized to an opaque
|
|
`INTERNAL_ERROR` so the in-process path matches the wire.
|
|
|
|
Returns:
|
|
A `(client, server)` pair. The wiring is symmetric, so the roles
|
|
are conventional only.
|
|
"""
|
|
ctx = TransportContext(kind=DIRECT_TRANSPORT_KIND, can_send_request=can_send_request, headers=headers)
|
|
client = DirectDispatcher(ctx, raise_handler_exceptions=raise_handler_exceptions)
|
|
server = DirectDispatcher(ctx, raise_handler_exceptions=raise_handler_exceptions)
|
|
client.connect_to(server)
|
|
server.connect_to(client)
|
|
return client, server
|