[v2] Dispatcher/ServerRunner receive-path swap — replaces BaseSession (#2710)
This commit is contained in:
Max
GitHub
parent
b478bff56d
commit
5d826490b6
25 changed files with 6594 additions and 1701 deletions
|
|
@ -14,6 +14,9 @@ from mcp.server import Server
|
|||
from mcp.server.mcpserver import MCPServer
|
||||
from mcp.shared.memory import create_client_server_memory_streams
|
||||
|
||||
SERVER_SHUTDOWN_GRACE = 2.0
|
||||
"""Seconds to wait for the in-process server to exit on EOF before cancelling."""
|
||||
|
||||
|
||||
class InMemoryTransport:
|
||||
"""In-memory transport for testing MCP servers without network overhead.
|
||||
|
|
@ -48,21 +51,49 @@ class InMemoryTransport:
|
|||
client_read, client_write = client_streams
|
||||
server_read, server_write = server_streams
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Start server in background
|
||||
tg.start_soon(
|
||||
lambda: actual_server.run(
|
||||
server_done = anyio.Event()
|
||||
|
||||
async def _run_server() -> None:
|
||||
try:
|
||||
await actual_server.run(
|
||||
server_read,
|
||||
server_write,
|
||||
actual_server.create_initialization_options(),
|
||||
raise_exceptions=self._raise_exceptions,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
server_done.set()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(_run_server)
|
||||
|
||||
try:
|
||||
yield client_read, client_write
|
||||
finally:
|
||||
tg.cancel_scope.cancel()
|
||||
# EOF the server (and our own read side) instead of
|
||||
# cancelling outright. The dispatcher's run() cancels its
|
||||
# own in-flight handlers on read-stream EOF, so for a
|
||||
# well-behaved server the task exits naturally and the
|
||||
# task-group join below is immediate. Cancelling here
|
||||
# unconditionally would `coro.throw()` into this task,
|
||||
# which on CPython 3.11 (gh-106749) drops `'call'` trace
|
||||
# events for the outer await chain and desyncs coverage's
|
||||
# CTracer past the test frame.
|
||||
await client_write.aclose()
|
||||
await server_write.aclose()
|
||||
# Backstop: the dispatcher exits on EOF, but the server's
|
||||
# own teardown (lifespan __aexit__, connection.exit_stack
|
||||
# callbacks) runs after that and is user code. If it never
|
||||
# completes the join would hang forever, so bound the wait
|
||||
# and fall back to cancelling. The healthy path returns
|
||||
# from wait() without the timeout firing, so the cancel is
|
||||
# never reached and gh-106749 stays avoided. If the cancel
|
||||
# does fire, the checkpoint at the end of
|
||||
# `create_client_server_memory_streams` resyncs the tracer.
|
||||
with anyio.move_on_after(SERVER_SHUTDOWN_GRACE):
|
||||
await server_done.wait()
|
||||
if not server_done.is_set():
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
async def __aenter__(self) -> TransportStreams:
|
||||
"""Connect to the server and return streams for communication."""
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from anyio.abc import TaskStatus
|
|||
from httpx_sse import SSEError, aconnect_sse
|
||||
|
||||
from mcp import types
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.shared._context_streams import create_context_streams
|
||||
from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
|
||||
from mcp.shared.message import SessionMessage
|
||||
|
|
@ -157,3 +158,4 @@ async def sse_client(
|
|||
|
||||
yield read_stream, write_stream
|
||||
tg.cancel_scope.cancel()
|
||||
await resync_tracer()
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
|
|||
from pydantic import ValidationError
|
||||
|
||||
from mcp.client._transport import TransportStreams
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
|
||||
from mcp.shared._httpx_utils import create_mcp_http_client
|
||||
from mcp.shared.message import ClientMessageMetadata, SessionMessage
|
||||
|
|
@ -586,3 +587,4 @@ async def streamable_http_client(
|
|||
if transport.session_id and terminate_on_close:
|
||||
await transport.terminate_session(client)
|
||||
tg.cancel_scope.cancel()
|
||||
await resync_tracer()
|
||||
|
|
|
|||
|
|
@ -1,14 +1,11 @@
|
|||
import importlib.metadata
|
||||
import logging
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
import anyio
|
||||
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.server.lowlevel.server import Server
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.types import ServerCapabilities
|
||||
|
||||
if not sys.warnoptions:
|
||||
warnings.simplefilter("ignore")
|
||||
|
|
@ -17,32 +14,10 @@ logging.basicConfig(level=logging.INFO)
|
|||
logger = logging.getLogger("server")
|
||||
|
||||
|
||||
async def receive_loop(session: ServerSession):
|
||||
logger.info("Starting receive loop")
|
||||
async for message in session.incoming_messages:
|
||||
if isinstance(message, Exception):
|
||||
logger.error("Error: %s", message)
|
||||
continue
|
||||
|
||||
logger.info("Received message from client: %s", message)
|
||||
|
||||
|
||||
async def main():
|
||||
version = importlib.metadata.version("mcp")
|
||||
async def main() -> None:
|
||||
server: Server[dict[str, object]] = Server("mcp")
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
async with (
|
||||
ServerSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
InitializationOptions(
|
||||
server_name="mcp",
|
||||
server_version=version,
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as session,
|
||||
write_stream,
|
||||
):
|
||||
await receive_loop(session)
|
||||
await server.run(read_stream, write_stream, server.create_initialization_options())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
249
src/mcp/server/connection.py
Normal file
249
src/mcp/server/connection.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
"""`Connection` - per-client connection state and the standalone outbound channel.
|
||||
|
||||
Always present on `Context` (never `None`), even in stateless deployments.
|
||||
Holds peer info populated at `initialize` time, per-connection scratch
|
||||
`state` and an `exit_stack` for teardown, and an `Outbound` for the
|
||||
standalone stream (the SSE GET stream in streamable HTTP, or the single duplex
|
||||
stream in stdio).
|
||||
|
||||
`notify` is best-effort: it never raises. If there's no standalone channel
|
||||
(stateless HTTP) or the stream has been dropped, the notification is
|
||||
debug-logged and silently discarded - server-initiated notifications are
|
||||
inherently advisory. `send_raw_request` *does* raise `NoBackChannelError` when
|
||||
there's no channel; `ping` is the only spec-sanctioned standalone request.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any, TypeVar, overload
|
||||
|
||||
import anyio
|
||||
from pydantic import BaseModel
|
||||
|
||||
from mcp.shared.dispatcher import CallOptions, Outbound
|
||||
from mcp.shared.exceptions import NoBackChannelError
|
||||
from mcp.shared.peer import Meta, dump_params
|
||||
from mcp.types import (
|
||||
ClientCapabilities,
|
||||
CreateMessageRequest,
|
||||
CreateMessageResult,
|
||||
ElicitRequest,
|
||||
ElicitResult,
|
||||
EmptyResult,
|
||||
InitializeRequestParams,
|
||||
ListRootsRequest,
|
||||
ListRootsResult,
|
||||
LoggingLevel,
|
||||
PingRequest,
|
||||
Request,
|
||||
)
|
||||
|
||||
__all__ = ["Connection"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ResultT = TypeVar("ResultT", bound=BaseModel)
|
||||
|
||||
# Result types for the spec's server-to-client request set, used by
|
||||
# `Connection.send_request` to infer the result type. If the spec's request
|
||||
# set grows substantially, consider declaring the result mapping on the
|
||||
# request types themselves (a `__mcp_result__` ClassVar read via a structural
|
||||
# protocol) so this table and the overload ladder don't need maintaining.
|
||||
_RESULT_FOR: dict[type[Request[Any, Any]], type[BaseModel]] = {
|
||||
CreateMessageRequest: CreateMessageResult,
|
||||
ElicitRequest: ElicitResult,
|
||||
ListRootsRequest: ListRootsResult,
|
||||
PingRequest: EmptyResult,
|
||||
}
|
||||
|
||||
|
||||
def _notification_params(payload: dict[str, Any] | None, meta: Meta | None) -> dict[str, Any] | None:
|
||||
if not meta:
|
||||
return payload
|
||||
out = dict(payload or {})
|
||||
out["_meta"] = meta
|
||||
return out
|
||||
|
||||
|
||||
class Connection:
|
||||
"""Per-client connection state and standalone-stream `Outbound`.
|
||||
|
||||
Constructed by `ServerRunner` once per connection. The peer-info fields
|
||||
are `None` until `initialize` completes; `initialized` is set later, when
|
||||
the client's `notifications/initialized` follow-up arrives. In stateless
|
||||
deployments the runner sets `initialized` immediately and peer-info
|
||||
remains `None` (no handshake reaches a stateless connection).
|
||||
"""
|
||||
|
||||
has_standalone_channel: bool
|
||||
session_id: str | None
|
||||
|
||||
client_params: InitializeRequestParams | None
|
||||
"""The full `initialize` request params; `None` before initialization."""
|
||||
|
||||
protocol_version: str | None
|
||||
|
||||
initialized: anyio.Event
|
||||
"""Set when `notifications/initialized` arrives (matches TS `oninitialized`);
|
||||
the point from which the spec permits server-initiated requests beyond
|
||||
ping/logging. Pre-set on stateless connections."""
|
||||
|
||||
state: dict[str, Any]
|
||||
"""Per-connection scratch state; persists across requests on this connection."""
|
||||
|
||||
exit_stack: AsyncExitStack
|
||||
"""Per-connection teardown, unwound LIFO (shielded) when the connection
|
||||
closes. Push cleanup from handlers or middleware; exceptions are logged
|
||||
and swallowed."""
|
||||
|
||||
def __init__(self, outbound: Outbound, *, has_standalone_channel: bool, session_id: str | None = None) -> None:
|
||||
self._outbound = outbound
|
||||
self.has_standalone_channel = has_standalone_channel
|
||||
self.session_id = session_id
|
||||
|
||||
self.client_params = None
|
||||
self.protocol_version = None
|
||||
self.initialized = anyio.Event()
|
||||
|
||||
self.state = {}
|
||||
|
||||
self.exit_stack = AsyncExitStack()
|
||||
|
||||
@property
|
||||
def initialize_accepted(self) -> bool:
|
||||
"""True once the inbound request gate is open: `initialize` recorded the
|
||||
peer info, or the handshake completed outright (stateless birth, or a
|
||||
bare `notifications/initialized`). Derived, never stored."""
|
||||
return self.client_params is not None or self.initialized.is_set()
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Send a raw request on the standalone stream.
|
||||
|
||||
Low-level `Outbound` channel. Prefer the typed `send_request` or the
|
||||
convenience methods below; use this directly only for off-spec
|
||||
messages. `opts` carries per-call `timeout` / `on_progress` /
|
||||
resumption hints; see `CallOptions`.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: `has_standalone_channel` is `False`.
|
||||
"""
|
||||
if not self.has_standalone_channel:
|
||||
raise NoBackChannelError(method)
|
||||
return await self._outbound.send_raw_request(method, params, opts)
|
||||
|
||||
@overload
|
||||
async def send_request(
|
||||
self, req: CreateMessageRequest, *, opts: CallOptions | None = None
|
||||
) -> CreateMessageResult: ...
|
||||
@overload
|
||||
async def send_request(self, req: ElicitRequest, *, opts: CallOptions | None = None) -> ElicitResult: ...
|
||||
@overload
|
||||
async def send_request(self, req: ListRootsRequest, *, opts: CallOptions | None = None) -> ListRootsResult: ...
|
||||
@overload
|
||||
async def send_request(self, req: PingRequest, *, opts: CallOptions | None = None) -> EmptyResult: ...
|
||||
@overload
|
||||
async def send_request(
|
||||
self, req: Request[Any, Any], *, result_type: type[ResultT], opts: CallOptions | None = None
|
||||
) -> ResultT: ...
|
||||
async def send_request(
|
||||
self,
|
||||
req: Request[Any, Any],
|
||||
*,
|
||||
result_type: type[BaseModel] | None = None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> BaseModel:
|
||||
"""Send a typed server-to-client request and return its typed result.
|
||||
|
||||
For spec request types the result type is inferred. For custom requests
|
||||
pass `result_type=` explicitly.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: No back-channel for server-initiated requests.
|
||||
pydantic.ValidationError: The peer's result does not match the expected result type.
|
||||
KeyError: `result_type` omitted for a non-spec request type.
|
||||
"""
|
||||
raw = await self.send_raw_request(req.method, dump_params(req.params), opts)
|
||||
cls = result_type if result_type is not None else _RESULT_FOR[type(req)]
|
||||
return cls.model_validate(raw, by_name=False)
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
"""Send a best-effort notification on the standalone stream.
|
||||
|
||||
Never raises. If there's no standalone channel or the stream is broken,
|
||||
the notification is dropped and debug-logged.
|
||||
"""
|
||||
if not self.has_standalone_channel:
|
||||
logger.debug("dropped %s: no standalone channel", method)
|
||||
return
|
||||
try:
|
||||
await self._outbound.notify(method, params)
|
||||
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
|
||||
logger.debug("dropped %s: standalone stream closed", method)
|
||||
|
||||
async def ping(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> None:
|
||||
"""Send a `ping` request on the standalone stream.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: `has_standalone_channel` is `False`.
|
||||
"""
|
||||
await self.send_raw_request("ping", dump_params(None, meta), opts)
|
||||
|
||||
async def log(self, level: LoggingLevel, data: Any, logger: str | None = None, *, meta: Meta | None = None) -> None:
|
||||
"""Send a `notifications/message` log entry on the standalone stream. Best-effort."""
|
||||
params: dict[str, Any] = {"level": level, "data": data}
|
||||
if logger is not None:
|
||||
params["logger"] = logger
|
||||
await self.notify("notifications/message", _notification_params(params, meta))
|
||||
|
||||
async def send_tool_list_changed(self, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/tools/list_changed", _notification_params(None, meta))
|
||||
|
||||
async def send_prompt_list_changed(self, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/prompts/list_changed", _notification_params(None, meta))
|
||||
|
||||
async def send_resource_list_changed(self, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/resources/list_changed", _notification_params(None, meta))
|
||||
|
||||
async def send_resource_updated(self, uri: str, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/resources/updated", _notification_params({"uri": uri}, meta))
|
||||
|
||||
def check_capability(self, capability: ClientCapabilities) -> bool:
|
||||
"""Return whether the connected client declared the given capability.
|
||||
|
||||
Returns `False` if `initialize` hasn't completed yet.
|
||||
"""
|
||||
# TODO: redesign - mirrors v1 ServerSession.check_client_capability
|
||||
# verbatim for parity.
|
||||
if self.client_params is None:
|
||||
return False
|
||||
have = self.client_params.capabilities
|
||||
if capability.roots is not None:
|
||||
if have.roots is None:
|
||||
return False
|
||||
if capability.roots.list_changed and not have.roots.list_changed:
|
||||
return False
|
||||
if capability.sampling is not None:
|
||||
if have.sampling is None:
|
||||
return False
|
||||
if capability.sampling.context is not None and have.sampling.context is None:
|
||||
return False
|
||||
if capability.sampling.tools is not None and have.sampling.tools is None:
|
||||
return False
|
||||
if capability.elicitation is not None and have.elicitation is None:
|
||||
return False
|
||||
if capability.experimental is not None:
|
||||
if have.experimental is None:
|
||||
return False
|
||||
for k, v in capability.experimental.items():
|
||||
if k not in have.experimental or have.experimental[k] != v:
|
||||
return False
|
||||
return True
|
||||
|
|
@ -1,21 +1,165 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic
|
||||
from typing import Any, Generic, Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared._context import RequestContext
|
||||
from mcp.shared.context import BaseContext
|
||||
from mcp.shared.dispatcher import DispatchContext
|
||||
from mcp.shared.message import CloseSSEStreamCallback
|
||||
from mcp.shared.peer import Meta
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import LoggingLevel, RequestId, RequestParamsMeta
|
||||
|
||||
# Invariant: parameterizes a mutable dataclass field; dict default matches the default lifespan.
|
||||
LifespanContextT = TypeVar("LifespanContextT", default=dict[str, Any])
|
||||
RequestT = TypeVar("RequestT", default=Any)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ServerRequestContext(RequestContext[ServerSession], Generic[LifespanContextT, RequestT]):
|
||||
class ServerRequestContext(Generic[LifespanContextT, RequestT]):
|
||||
"""Per-request context handed to lowlevel request and notification handlers.
|
||||
|
||||
Built by `ServerRunner._make_context` for each inbound message. Carries the
|
||||
connection-scoped `ServerSession` (server-to-client requests and
|
||||
notifications), per-request metadata, and any per-message data the
|
||||
transport attached (the HTTP request, SSE stream-close callbacks).
|
||||
"""
|
||||
|
||||
session: ServerSession
|
||||
lifespan_context: LifespanContextT
|
||||
request_id: RequestId | None = None
|
||||
meta: RequestParamsMeta | None = None
|
||||
request: RequestT | None = None
|
||||
close_sse_stream: CloseSSEStreamCallback | None = None
|
||||
close_standalone_sse_stream: CloseSSEStreamCallback | None = None
|
||||
|
||||
|
||||
# Covariant: `lifespan` is exposed read-only, so a `Context[AppState]` passes as `Context[object]`.
|
||||
LifespanT_co = TypeVar("LifespanT_co", default=Any, covariant=True)
|
||||
|
||||
|
||||
class Context(BaseContext[TransportContext], Generic[LifespanT_co]):
|
||||
"""Server-side per-request context.
|
||||
|
||||
Extends `BaseContext` (transport metadata, the raw back-channel, progress
|
||||
reporting) with `lifespan`, `connection`, and request-scoped `log`.
|
||||
|
||||
Not currently constructed by `ServerRunner`, which hands handlers a
|
||||
`ServerRequestContext` instead.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
*,
|
||||
lifespan: LifespanT_co,
|
||||
connection: Connection,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
) -> None:
|
||||
super().__init__(dctx, meta=meta)
|
||||
self._lifespan = lifespan
|
||||
self._connection = connection
|
||||
|
||||
@property
|
||||
def lifespan(self) -> LifespanT_co:
|
||||
"""The server-wide lifespan output (what `Server(..., lifespan=...)` yielded)."""
|
||||
return self._lifespan
|
||||
|
||||
@property
|
||||
def connection(self) -> Connection:
|
||||
"""The per-client `Connection` for this request's connection."""
|
||||
return self._connection
|
||||
|
||||
@property
|
||||
def session_id(self) -> str | None:
|
||||
"""The transport's session id for this connection, when one exists.
|
||||
|
||||
Convenience for `ctx.connection.session_id`. `None` on stdio and
|
||||
stateless HTTP.
|
||||
"""
|
||||
return self._connection.session_id
|
||||
|
||||
@property
|
||||
def headers(self) -> Mapping[str, str] | None:
|
||||
"""Request headers carried by this message, when the transport has them.
|
||||
|
||||
Convenience for `ctx.transport.headers`. `None` on stdio.
|
||||
"""
|
||||
return self.transport.headers
|
||||
|
||||
async def log(self, level: LoggingLevel, data: Any, logger: str | None = None, *, meta: Meta | None = None) -> None:
|
||||
"""Send a request-scoped `notifications/message` log entry.
|
||||
|
||||
Uses this request's back-channel (so the entry rides the request's SSE
|
||||
stream in streamable HTTP), not the standalone stream - use
|
||||
`ctx.connection.log(...)` for that.
|
||||
"""
|
||||
params: dict[str, Any] = {"level": level, "data": data}
|
||||
if logger is not None:
|
||||
params["logger"] = logger
|
||||
if meta:
|
||||
params["_meta"] = meta
|
||||
await self.notify("notifications/message", params)
|
||||
|
||||
|
||||
HandlerResult = BaseModel | dict[str, Any] | None
|
||||
"""What a request handler (or middleware) may return. `ServerRunner` serializes
|
||||
all three to a result dict."""
|
||||
|
||||
CallNext = Callable[[], Awaitable[HandlerResult]]
|
||||
|
||||
_MwLifespanT = TypeVar("_MwLifespanT")
|
||||
|
||||
|
||||
class ServerMiddleware(Protocol[_MwLifespanT]):
|
||||
"""Context-tier middleware: `(ctx, method, params, call_next) -> result`.
|
||||
|
||||
Runs at the top of `ServerRunner._on_request` / `_on_notify` after `ctx`
|
||||
is built but before any validation, lookup, or handshake. Wraps every
|
||||
inbound request and notification: `initialize`, the pre-init gate,
|
||||
`METHOD_NOT_FOUND`, params validation, the handler call, and
|
||||
`notifications/initialized` all run inside `call_next()`.
|
||||
`notifications/cancelled` is observed too; the dispatcher applies the
|
||||
cancellation itself, then forwards the notification. A request-side
|
||||
failure reaches the middleware as a raised `MCPError` (or
|
||||
`ValidationError` for malformed params) so observation/logging middleware
|
||||
can record it. Listed outermost-first on `Server.middleware`.
|
||||
|
||||
`ctx.request_id is None` distinguishes a notification from a request. For
|
||||
notifications `call_next()` returns `None` (a dropped or unhandled
|
||||
notification also returns `None`) and the middleware's own return value is
|
||||
discarded.
|
||||
|
||||
`params` is the raw inbound mapping (no model validation has happened
|
||||
yet). For typed inspection, validate against the model the middleware
|
||||
expects.
|
||||
|
||||
Warning: `initialize` is handled inline - the dispatcher does not read
|
||||
further inbound messages until the middleware chain returns. Awaiting a
|
||||
server-to-client request (`ctx.session.send_request`, `send_ping`, ...)
|
||||
while handling `initialize` therefore deadlocks the connection: the
|
||||
response can never be dequeued. Send-and-forget notifications are safe.
|
||||
|
||||
`Server[L].middleware` holds `ServerMiddleware[L]`, so an app-specific
|
||||
middleware sees `ctx.lifespan_context: L`. While the context is the
|
||||
mutable `ServerRequestContext` dataclass it is invariant in `L`, so a
|
||||
reusable middleware should be typed `ServerMiddleware[Any]` to register on
|
||||
any `Server[L]`.
|
||||
"""
|
||||
|
||||
# TODO(maxisbey): once `_make_context` returns the (covariant) `Context[L]`
|
||||
# again, restore `_MwLifespanT` to `contravariant=True` and retype `ctx`
|
||||
# below to `Context[_MwLifespanT]` so reusable middleware can be
|
||||
# `ServerMiddleware[object]` instead of `ServerMiddleware[Any]`.
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
ctx: ServerRequestContext[_MwLifespanT, Any],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
call_next: CallNext,
|
||||
) -> HandlerResult: ...
|
||||
|
|
|
|||
|
|
@ -36,16 +36,14 @@ handler callables by method string.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
import logging
|
||||
import warnings
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from importlib.metadata import version as importlib_version
|
||||
from typing import Any, Generic, cast
|
||||
from typing import Any, Generic
|
||||
|
||||
import anyio
|
||||
from opentelemetry.trace import SpanKind, StatusCode
|
||||
from pydantic import BaseModel
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.authentication import AuthenticationMiddleware
|
||||
|
|
@ -58,22 +56,45 @@ from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend, RequireAut
|
|||
from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenVerifier
|
||||
from mcp.server.auth.routes import build_resource_metadata_url, create_auth_routes, create_protected_resource_routes
|
||||
from mcp.server.auth.settings import AuthSettings
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.server.context import HandlerResult, ServerMiddleware, ServerRequestContext
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.server.runner import ServerRunner, otel_middleware
|
||||
from mcp.server.streamable_http import EventStore
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared._otel import extract_trace_context, otel_span
|
||||
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.message import ServerMessageMetadata, SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LifespanResultT = TypeVar("LifespanResultT", default=Any)
|
||||
|
||||
_ParamsT = TypeVar("_ParamsT", bound=BaseModel, default=BaseModel)
|
||||
|
||||
RequestHandler = Callable[[ServerRequestContext[LifespanResultT], _ParamsT], Awaitable[HandlerResult]]
|
||||
"""A registered request handler: `(ctx, params) -> result`."""
|
||||
|
||||
NotificationHandler = Callable[[ServerRequestContext[LifespanResultT], _ParamsT], Awaitable[None]]
|
||||
"""A registered notification handler: `(ctx, params) -> None`."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HandlerEntry(Generic[LifespanResultT]):
|
||||
"""A registered handler and the params model to validate incoming params against.
|
||||
|
||||
Stored in `Server._request_handlers` / `_notification_handlers` and consumed
|
||||
by `ServerRunner` to validate, build `Context`, and invoke. The handler's
|
||||
second-argument type is erased to `Any` in storage (each entry has a
|
||||
different concrete params type and `Callable` parameters are contravariant);
|
||||
the precise type is recoverable via `params_type`. The correlation is
|
||||
enforced at registration time by `Server.add_request_handler`.
|
||||
"""
|
||||
|
||||
params_type: type[BaseModel]
|
||||
handler: RequestHandler[LifespanResultT, Any]
|
||||
|
||||
|
||||
class NotificationOptions:
|
||||
def __init__(self, prompts_changed: bool = False, resources_changed: bool = False, tools_changed: bool = False):
|
||||
|
|
@ -83,7 +104,7 @@ class NotificationOptions:
|
|||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: Server[LifespanResultT]) -> AsyncIterator[dict[str, Any]]:
|
||||
async def lifespan(_: Server[Any]) -> AsyncIterator[dict[str, Any]]:
|
||||
"""Default lifespan context manager that does nothing.
|
||||
|
||||
Returns:
|
||||
|
|
@ -191,53 +212,91 @@ class Server(Generic[LifespanResultT]):
|
|||
self.website_url = website_url
|
||||
self.icons = icons
|
||||
self.lifespan = lifespan
|
||||
self._request_handlers: dict[str, Callable[[ServerRequestContext[LifespanResultT], Any], Awaitable[Any]]] = {}
|
||||
self._notification_handlers: dict[
|
||||
str, Callable[[ServerRequestContext[LifespanResultT], Any], Awaitable[None]]
|
||||
] = {}
|
||||
self._request_handlers: dict[str, HandlerEntry[LifespanResultT]] = {}
|
||||
self._notification_handlers: dict[str, HandlerEntry[LifespanResultT]] = {}
|
||||
self._session_manager: StreamableHTTPSessionManager | None = None
|
||||
# Context-tier middleware: wraps every inbound request (including
|
||||
# `initialize`, lookup, validation, handler) with
|
||||
# `(ctx, method, params, call_next)`. Applied in `ServerRunner._on_request`.
|
||||
# TODO(maxisbey): provisional - signature and semantics change with the
|
||||
# Context/middleware rework (covariant `Context[L]`, outbound seam) before
|
||||
# v2 final.
|
||||
self.middleware: list[ServerMiddleware[LifespanResultT]] = []
|
||||
logger.debug("Initializing server %r", name)
|
||||
|
||||
# Populate internal handler dicts from on_* kwargs
|
||||
self._request_handlers.update(
|
||||
{
|
||||
method: handler
|
||||
for method, handler in {
|
||||
"ping": on_ping,
|
||||
"prompts/list": on_list_prompts,
|
||||
"prompts/get": on_get_prompt,
|
||||
"resources/list": on_list_resources,
|
||||
"resources/templates/list": on_list_resource_templates,
|
||||
"resources/read": on_read_resource,
|
||||
"resources/subscribe": on_subscribe_resource,
|
||||
"resources/unsubscribe": on_unsubscribe_resource,
|
||||
"tools/list": on_list_tools,
|
||||
"tools/call": on_call_tool,
|
||||
"logging/setLevel": on_set_logging_level,
|
||||
"completion/complete": on_completion,
|
||||
}.items()
|
||||
if handler is not None
|
||||
}
|
||||
)
|
||||
_spec_requests: list[tuple[str, type[BaseModel], RequestHandler[LifespanResultT, Any] | None]] = [
|
||||
("ping", types.RequestParams, on_ping),
|
||||
("prompts/list", types.PaginatedRequestParams, on_list_prompts),
|
||||
("prompts/get", types.GetPromptRequestParams, on_get_prompt),
|
||||
("resources/list", types.PaginatedRequestParams, on_list_resources),
|
||||
("resources/templates/list", types.PaginatedRequestParams, on_list_resource_templates),
|
||||
("resources/read", types.ReadResourceRequestParams, on_read_resource),
|
||||
("resources/subscribe", types.SubscribeRequestParams, on_subscribe_resource),
|
||||
("resources/unsubscribe", types.UnsubscribeRequestParams, on_unsubscribe_resource),
|
||||
("tools/list", types.PaginatedRequestParams, on_list_tools),
|
||||
("tools/call", types.CallToolRequestParams, on_call_tool),
|
||||
("logging/setLevel", types.SetLevelRequestParams, on_set_logging_level),
|
||||
("completion/complete", types.CompleteRequestParams, on_completion),
|
||||
]
|
||||
self._request_handlers.update({m: HandlerEntry(pt, h) for m, pt, h in _spec_requests if h is not None})
|
||||
|
||||
_spec_notifications: list[tuple[str, type[BaseModel], NotificationHandler[LifespanResultT, Any] | None]] = [
|
||||
("notifications/roots/list_changed", types.NotificationParams, on_roots_list_changed),
|
||||
("notifications/progress", types.ProgressNotificationParams, on_progress),
|
||||
]
|
||||
self._notification_handlers.update(
|
||||
{
|
||||
method: handler
|
||||
for method, handler in {
|
||||
"notifications/roots/list_changed": on_roots_list_changed,
|
||||
"notifications/progress": on_progress,
|
||||
}.items()
|
||||
if handler is not None
|
||||
}
|
||||
{m: HandlerEntry(pt, h) for m, pt, h in _spec_notifications if h is not None}
|
||||
)
|
||||
|
||||
def _add_request_handler(
|
||||
def add_request_handler(
|
||||
self,
|
||||
method: str,
|
||||
handler: Callable[[ServerRequestContext[LifespanResultT], Any], Awaitable[Any]],
|
||||
params_type: type[_ParamsT],
|
||||
handler: RequestHandler[LifespanResultT, _ParamsT],
|
||||
) -> None:
|
||||
"""Add a request handler, silently replacing any existing handler for the same method."""
|
||||
self._request_handlers[method] = handler
|
||||
"""Register a request handler for `method`.
|
||||
|
||||
`params_type` is the model incoming params are validated against
|
||||
before the handler is invoked. It should subclass `RequestParams` so
|
||||
`_meta` parses uniformly. A message with no `params` member validates
|
||||
`{}` against `params_type`: models with required fields reject it as
|
||||
INVALID_PARAMS, all-optional models reach the handler with their
|
||||
defaults - the handler never receives `None`. Replaces any existing
|
||||
handler for the same method, except `initialize`, which is reserved:
|
||||
the runner owns the handshake, so registering it raises `ValueError`.
|
||||
Use `Server.middleware` to observe or wrap initialization.
|
||||
"""
|
||||
if method == "initialize":
|
||||
raise ValueError(
|
||||
"'initialize' is handled by the server runner and cannot be overridden; "
|
||||
"use Server.middleware to observe or wrap initialization"
|
||||
)
|
||||
self._request_handlers[method] = HandlerEntry(params_type, handler)
|
||||
|
||||
def add_notification_handler(
|
||||
self,
|
||||
method: str,
|
||||
params_type: type[_ParamsT],
|
||||
handler: NotificationHandler[LifespanResultT, _ParamsT],
|
||||
) -> None:
|
||||
"""Register a notification handler for `method`.
|
||||
|
||||
`params_type` should subclass `NotificationParams` so `_meta`
|
||||
parses uniformly. Absent params follow the same contract as requests:
|
||||
`{}` is validated, so the handler receives the model with its defaults,
|
||||
never `None`. Replaces any existing handler. A handler for
|
||||
`notifications/initialized` runs after the runner has marked the
|
||||
connection initialized.
|
||||
"""
|
||||
self._notification_handlers[method] = HandlerEntry(params_type, handler)
|
||||
|
||||
def get_request_handler(self, method: str) -> HandlerEntry[LifespanResultT] | None:
|
||||
"""Return the registered entry for a request method, or `None`."""
|
||||
return self._request_handlers.get(method)
|
||||
|
||||
def get_notification_handler(self, method: str) -> HandlerEntry[LifespanResultT] | None:
|
||||
"""Return the registered entry for a notification method, or `None`."""
|
||||
return self._notification_handlers.get(method)
|
||||
|
||||
# TODO: Rethink capabilities API. Currently capabilities are derived from registered
|
||||
# handlers but require NotificationOptions to be passed externally for list_changed
|
||||
|
|
@ -347,167 +406,31 @@ class Server(Generic[LifespanResultT]):
|
|||
# the initialization lifecycle, but can do so with any available node
|
||||
# rather than requiring initialization for each connection.
|
||||
stateless: bool = False,
|
||||
):
|
||||
async with AsyncExitStack() as stack:
|
||||
lifespan_context = await stack.enter_async_context(self.lifespan(self))
|
||||
session = await stack.enter_async_context(
|
||||
ServerSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
initialization_options,
|
||||
stateless=stateless,
|
||||
)
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
try:
|
||||
async for message in session.incoming_messages:
|
||||
logger.debug("Received message: %s", message)
|
||||
|
||||
if isinstance(message, RequestResponder) and message.context is not None:
|
||||
context = message.context
|
||||
else:
|
||||
context = contextvars.copy_context()
|
||||
|
||||
context.run(
|
||||
tg.start_soon,
|
||||
self._handle_message,
|
||||
message,
|
||||
session,
|
||||
lifespan_context,
|
||||
raise_exceptions,
|
||||
)
|
||||
finally:
|
||||
# Transport closed: cancel in-flight handlers. Without this the
|
||||
# TG join waits for them, and when they eventually try to
|
||||
# respond they hit a closed write stream (the session's
|
||||
# _receive_loop closed it when the read stream ended).
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
async def _handle_message(
|
||||
self,
|
||||
message: RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
raise_exceptions: bool = False,
|
||||
):
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
match message:
|
||||
case RequestResponder() as responder:
|
||||
with responder:
|
||||
await self._handle_request(
|
||||
message, responder.request, session, lifespan_context, raise_exceptions
|
||||
)
|
||||
case Exception():
|
||||
logger.error(f"Received exception from stream: {message}")
|
||||
if raise_exceptions:
|
||||
raise message
|
||||
case _:
|
||||
await self._handle_notification(message, session, lifespan_context)
|
||||
|
||||
for warning in w: # pragma: lax no cover
|
||||
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
|
||||
|
||||
async def _handle_request(
|
||||
self,
|
||||
message: RequestResponder[types.ClientRequest, types.ServerResult],
|
||||
req: types.ClientRequest,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
raise_exceptions: bool,
|
||||
):
|
||||
logger.info("Processing request of type %s", type(req).__name__)
|
||||
|
||||
target = getattr(req.params, "name", None) if req.params else None
|
||||
span_name = f"MCP handle {req.method} {target}" if target else f"MCP handle {req.method}"
|
||||
|
||||
# Extract W3C trace context from _meta (SEP-414).
|
||||
meta = cast(dict[str, Any] | None, getattr(req.params, "meta", None)) if req.params else None
|
||||
parent_context = extract_trace_context(meta) if meta is not None else None
|
||||
|
||||
with otel_span(
|
||||
span_name,
|
||||
kind=SpanKind.SERVER,
|
||||
attributes={"mcp.method.name": req.method, "jsonrpc.request.id": message.request_id},
|
||||
context=parent_context,
|
||||
) as span:
|
||||
if handler := self._request_handlers.get(req.method):
|
||||
logger.debug("Dispatching request of type %s", type(req).__name__)
|
||||
|
||||
try:
|
||||
# Extract request context and close_sse_stream from message metadata
|
||||
request_data = None
|
||||
close_sse_stream_cb = None
|
||||
close_standalone_sse_stream_cb = None
|
||||
if message.message_metadata is not None and isinstance(
|
||||
message.message_metadata, ServerMessageMetadata
|
||||
):
|
||||
request_data = message.message_metadata.request_context
|
||||
close_sse_stream_cb = message.message_metadata.close_sse_stream
|
||||
close_standalone_sse_stream_cb = message.message_metadata.close_standalone_sse_stream
|
||||
|
||||
ctx = ServerRequestContext(
|
||||
request_id=message.request_id,
|
||||
meta=message.request_meta,
|
||||
session=session,
|
||||
lifespan_context=lifespan_context,
|
||||
request=request_data,
|
||||
close_sse_stream=close_sse_stream_cb,
|
||||
close_standalone_sse_stream=close_standalone_sse_stream_cb,
|
||||
)
|
||||
response = await handler(ctx, req.params)
|
||||
except MCPError as err:
|
||||
response = err.error
|
||||
except anyio.get_cancelled_exc_class():
|
||||
if message.cancelled:
|
||||
# Client sent CancelledNotification; responder.cancel() already
|
||||
# sent an error response, so skip the duplicate.
|
||||
logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
|
||||
return
|
||||
# Transport-close cancellation from the TG in run(); re-raise so the
|
||||
# TG swallows its own cancellation.
|
||||
raise
|
||||
except Exception as err:
|
||||
if raise_exceptions: # pragma: no cover
|
||||
raise err
|
||||
response = types.ErrorData(code=0, message=str(err))
|
||||
else:
|
||||
response = types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found")
|
||||
|
||||
if isinstance(response, types.ErrorData) and span is not None:
|
||||
span.set_status(StatusCode.ERROR, response.message)
|
||||
|
||||
try:
|
||||
await message.respond(response)
|
||||
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
|
||||
# Transport closed between handler unblocking and respond. Happens
|
||||
# when _receive_loop's finally wakes a handler blocked on
|
||||
# send_request: the handler runs to respond() before run()'s TG
|
||||
# cancel fires, but after the write stream closed. Closed if our
|
||||
# end closed (_receive_loop's async-with exit); Broken if the peer
|
||||
# end closed first (streamable_http terminate()).
|
||||
logger.debug("Response for %s dropped - transport closed", message.request_id)
|
||||
return
|
||||
|
||||
logger.debug("Response sent")
|
||||
|
||||
async def _handle_notification(
|
||||
self,
|
||||
notify: types.ClientNotification,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
) -> None:
|
||||
if handler := self._notification_handlers.get(notify.method):
|
||||
logger.debug("Dispatching notification of type %s", type(notify).__name__)
|
||||
|
||||
try:
|
||||
ctx = ServerRequestContext(
|
||||
session=session,
|
||||
lifespan_context=lifespan_context,
|
||||
)
|
||||
await handler(ctx, notify.params)
|
||||
except Exception: # pragma: no cover
|
||||
logger.exception("Uncaught exception in notification handler")
|
||||
async with self.lifespan(self) as lifespan_context:
|
||||
dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
|
||||
read_stream,
|
||||
write_stream,
|
||||
raise_handler_exceptions=raise_exceptions,
|
||||
# Handle `initialize` inline so a client that pipelines it with
|
||||
# the next request (spec says SHOULD NOT, not MUST NOT) sees
|
||||
# the initialized state instead of failing the init-gate.
|
||||
inline_methods=frozenset({"initialize"}),
|
||||
)
|
||||
runner = ServerRunner(
|
||||
server=self,
|
||||
dispatcher=dispatcher,
|
||||
lifespan_state=lifespan_context,
|
||||
init_options=initialization_options,
|
||||
# Stateless HTTP has no standalone GET stream, so server-initiated
|
||||
# requests on `runner.connection` must fail fast with
|
||||
# `NoBackChannelError` rather than write to a channel that will
|
||||
# never deliver a response.
|
||||
has_standalone_channel=not stateless,
|
||||
stateless=stateless,
|
||||
dispatch_middleware=[otel_middleware],
|
||||
)
|
||||
await runner.run()
|
||||
|
||||
def streamable_http_app(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -607,11 +607,7 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
completion=result if result is not None else Completion(values=[], total=None, has_more=None),
|
||||
)
|
||||
|
||||
# TODO(maxisbey): remove private access — completion needs post-construction
|
||||
# handler registration, find a better pattern for this
|
||||
self._lowlevel_server._add_request_handler( # pyright: ignore[reportPrivateUsage]
|
||||
"completion/complete", handler
|
||||
)
|
||||
self._lowlevel_server.add_request_handler("completion/complete", CompleteRequestParams, handler)
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
|
|
|||
377
src/mcp/server/runner.py
Normal file
377
src/mcp/server/runner.py
Normal file
|
|
@ -0,0 +1,377 @@
|
|||
"""`ServerRunner` - per-connection orchestrator over a `Dispatcher`.
|
||||
|
||||
`ServerRunner` is the bridge between the dispatcher layer (`on_request` /
|
||||
`on_notify`, untyped dicts) and the user's handler layer (typed `Context`,
|
||||
typed params). One instance per client connection. It:
|
||||
|
||||
* handles the `initialize` handshake and populates `Connection`
|
||||
* gates requests until initialized (`ping` exempt)
|
||||
* looks up the handler in the server's registry, validates params, builds
|
||||
`Context`, runs the middleware chain, returns the result dict
|
||||
* drives `dispatcher.run()` and the per-connection lifespan
|
||||
|
||||
`ServerRunner` holds a `Server` directly - `Server` is the registry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial, reduce
|
||||
from typing import TYPE_CHECKING, Any, Generic, cast, get_args
|
||||
|
||||
import anyio.abc
|
||||
from opentelemetry.trace import SpanKind, StatusCode
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared._otel import extract_trace_context, otel_span
|
||||
from mcp.shared.dispatcher import DispatchContext, DispatchMiddleware, OnRequest
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.message import ServerMessageMetadata
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||
from mcp.types import (
|
||||
INVALID_PARAMS,
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
METHOD_NOT_FOUND,
|
||||
ClientRequest,
|
||||
ErrorData,
|
||||
Implementation,
|
||||
InitializeRequestParams,
|
||||
InitializeResult,
|
||||
NotificationParams,
|
||||
RequestParams,
|
||||
RequestParamsMeta,
|
||||
client_request_adapter,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.server.lowlevel.server import Server
|
||||
|
||||
__all__ = ["CallNext", "ServerMiddleware", "ServerRunner", "otel_middleware"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LifespanT = TypeVar("LifespanT", default=Any)
|
||||
|
||||
|
||||
_INIT_EXEMPT: frozenset[str] = frozenset({"ping"})
|
||||
|
||||
_EXIT_STACK_CLOSE_TIMEOUT: float = 5
|
||||
"""Bound for the shielded exit-stack unwind in `run()`; a hung cleanup
|
||||
callback must not wedge shutdown."""
|
||||
|
||||
|
||||
def _extract_meta(params: Mapping[str, Any] | None) -> RequestParamsMeta | None:
|
||||
"""Lift `_meta` from raw params; `None` when absent or malformed, so
|
||||
context construction is independent of params validity."""
|
||||
if not params or "_meta" not in params:
|
||||
return None
|
||||
try:
|
||||
return RequestParams.model_validate(params, by_name=False).meta
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
_SPEC_CLIENT_METHODS: frozenset[str] = frozenset(
|
||||
cast(type[BaseModel], arm).model_fields["method"].default for arm in get_args(ClientRequest)
|
||||
)
|
||||
"""Method names in the spec `ClientRequest` union, derived from the
|
||||
discriminator literal on each arm. Used to gate upfront validation so custom
|
||||
methods registered via `add_request_handler` are not rejected."""
|
||||
|
||||
|
||||
def otel_middleware(next_on_request: OnRequest) -> OnRequest:
|
||||
"""Dispatch-tier middleware that wraps each request in an OpenTelemetry span.
|
||||
|
||||
Mirrors the span shape of the existing `Server._handle_request`: span name
|
||||
`"MCP handle <method> [<target>]"`, `mcp.method.name` attribute, W3C
|
||||
trace context extracted from `params._meta` (SEP-414), and an ERROR
|
||||
status if the handler raises.
|
||||
"""
|
||||
|
||||
async def wrapped(
|
||||
dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
|
||||
) -> dict[str, Any]:
|
||||
target: str | None
|
||||
match params:
|
||||
case {"name": str() as target}:
|
||||
pass
|
||||
case _:
|
||||
target = None
|
||||
parent: Any | None
|
||||
match params:
|
||||
case {"_meta": {**meta}}:
|
||||
parent = extract_trace_context(meta)
|
||||
case _:
|
||||
parent = None
|
||||
span_name = f"MCP handle {method}{f' {target}' if target else ''}"
|
||||
# `otel_middleware` wraps `on_request` only, so `request_id` is always set.
|
||||
attributes = {"mcp.method.name": method, "jsonrpc.request.id": str(dctx.request_id)}
|
||||
with otel_span(
|
||||
span_name,
|
||||
kind=SpanKind.SERVER,
|
||||
attributes=attributes,
|
||||
context=parent,
|
||||
record_exception=False,
|
||||
set_status_on_exception=False,
|
||||
) as span:
|
||||
try:
|
||||
return await next_on_request(dctx, method, params)
|
||||
except MCPError as e:
|
||||
span.set_status(StatusCode.ERROR, e.error.message)
|
||||
raise
|
||||
except ValidationError:
|
||||
# Mirror the sanitized wire response; pydantic messages carry client input.
|
||||
span.set_status(StatusCode.ERROR, "Invalid request parameters")
|
||||
raise
|
||||
except Exception as e:
|
||||
span.record_exception(e)
|
||||
span.set_status(StatusCode.ERROR, str(e))
|
||||
raise
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _dump_result(result: Any) -> dict[str, Any]:
|
||||
if result is None:
|
||||
return {}
|
||||
if isinstance(result, ErrorData):
|
||||
# ErrorData is a JSON-RPC error, not a success result. Handler returns
|
||||
# already raise in `_inner`; this catches middleware returning one.
|
||||
raise MCPError.from_error_data(result)
|
||||
if isinstance(result, BaseModel):
|
||||
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
if isinstance(result, dict):
|
||||
return cast(dict[str, Any], result)
|
||||
raise TypeError(f"handler returned {type(result).__name__}; expected BaseModel, dict, or None")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerRunner(Generic[LifespanT]):
|
||||
"""Per-connection orchestrator. One instance per client connection."""
|
||||
|
||||
server: Server[LifespanT]
|
||||
dispatcher: JSONRPCDispatcher[Any]
|
||||
lifespan_state: LifespanT
|
||||
has_standalone_channel: bool
|
||||
init_options: InitializationOptions | None = None
|
||||
"""`InitializeResult` payload. Defaults to `server.create_initialization_options()`."""
|
||||
session_id: str | None = None
|
||||
stateless: bool = False
|
||||
dispatch_middleware: list[DispatchMiddleware] = field(default_factory=list[DispatchMiddleware])
|
||||
|
||||
connection: Connection = field(init=False)
|
||||
session: ServerSession = field(init=False)
|
||||
"""Connection-scoped: the same instance reaches every request as `ctx.session`."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.init_options is None:
|
||||
self.init_options = self.server.create_initialization_options()
|
||||
self.connection = Connection(
|
||||
self.dispatcher, has_standalone_channel=self.has_standalone_channel, session_id=self.session_id
|
||||
)
|
||||
if self.stateless:
|
||||
# No handshake ever arrives on a stateless connection; born ready.
|
||||
self.connection.initialized.set()
|
||||
self.session = ServerSession(self.dispatcher, self.connection, stateless=self.stateless)
|
||||
|
||||
async def run(self, *, task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED) -> None:
|
||||
"""Drive the dispatcher until the underlying channel closes.
|
||||
|
||||
Composes `dispatch_middleware` over `_on_request` and hands the result
|
||||
to `dispatcher.run()`. `task_status.started()` is forwarded so callers
|
||||
can `await tg.start(runner.run)` and resume once the dispatcher is
|
||||
ready to accept requests. Once the dispatcher exits,
|
||||
`connection.exit_stack` is unwound (shielded from outer cancellation,
|
||||
bounded by `_EXIT_STACK_CLOSE_TIMEOUT`) so any per-connection cleanup
|
||||
registered by handlers or middleware gets a chance to run without a
|
||||
misbehaving callback hanging shutdown indefinitely.
|
||||
"""
|
||||
try:
|
||||
await self.dispatcher.run(self._compose_on_request(), self._on_notify, task_status=task_status)
|
||||
finally:
|
||||
with anyio.move_on_after(_EXIT_STACK_CLOSE_TIMEOUT, shield=True) as scope:
|
||||
try:
|
||||
await self.connection.exit_stack.aclose()
|
||||
except Exception:
|
||||
# Raising here would mask dispatcher.run()'s exception and
|
||||
# crash stdio servers on normal disconnect.
|
||||
logger.exception("connection exit_stack cleanup raised")
|
||||
if scope.cancelled_caught:
|
||||
logger.warning(
|
||||
"connection exit_stack cleanup exceeded %s seconds; abandoning remaining callbacks",
|
||||
_EXIT_STACK_CLOSE_TIMEOUT,
|
||||
)
|
||||
|
||||
def _compose_on_request(self) -> OnRequest:
|
||||
"""Wrap `_on_request` in `dispatch_middleware`, outermost-first.
|
||||
|
||||
Dispatch-tier middleware sees raw `(dctx, method, params) -> dict`
|
||||
and wraps everything - initialize, METHOD_NOT_FOUND, validation
|
||||
failures included. `run()` calls this once and hands the result to
|
||||
`dispatcher.run()`.
|
||||
"""
|
||||
return reduce(lambda h, mw: mw(h), reversed(self.dispatch_middleware), self._on_request)
|
||||
|
||||
async def _on_request(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
ctx = self._make_context(dctx, _extract_meta(params))
|
||||
|
||||
async def _inner() -> HandlerResult:
|
||||
# TODO(maxisbey): pinned compat: spec methods are validated against
|
||||
# the ClientRequest union before lookup, so malformed params are
|
||||
# INVALID_PARAMS even with no handler registered.
|
||||
if method in _SPEC_CLIENT_METHODS:
|
||||
payload: dict[str, Any] = {"method": method}
|
||||
if params is not None:
|
||||
payload["params"] = dict(params)
|
||||
client_request_adapter.validate_python(payload, by_name=False)
|
||||
# TODO(maxisbey): the 2026-07-28 spec drops the handshake; this branch and
|
||||
# the gate become a per-version legacy path then. Initialize runs inline
|
||||
# (read loop parked), so awaiting the peer anywhere on this path deadlocks.
|
||||
if method == "initialize":
|
||||
return self._handle_initialize(params)
|
||||
if not self.connection.initialize_accepted and method not in _INIT_EXEMPT:
|
||||
# Pinned compat: the same error shape the union validation produced.
|
||||
raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="")
|
||||
entry = self.server.get_request_handler(method)
|
||||
if entry is None:
|
||||
raise MCPError(code=METHOD_NOT_FOUND, message="Method not found")
|
||||
# Absent params validate as {} (required fields still reject), so
|
||||
# the handler receives the model with its defaults, never None.
|
||||
typed_params = entry.params_type.model_validate({} if params is None else params, by_name=False)
|
||||
result = await entry.handler(ctx, typed_params)
|
||||
if isinstance(result, ErrorData):
|
||||
# Raise inside the chain so middleware observes the failure.
|
||||
raise MCPError.from_error_data(result)
|
||||
return result
|
||||
|
||||
call = self._compose_server_middleware(ctx, method, params, _inner)
|
||||
result = _dump_result(await call())
|
||||
if method == "initialize":
|
||||
# Commit only on chain success, so a middleware veto leaves no state.
|
||||
# Race-free: the read loop is parked until this call returns.
|
||||
self.connection.client_params, self.connection.protocol_version = self._negotiate_initialize(params)
|
||||
return result
|
||||
|
||||
async def _on_notify(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
) -> None:
|
||||
ctx = self._make_context(dctx, _extract_meta(params))
|
||||
|
||||
async def _inner() -> None:
|
||||
if method == "notifications/initialized":
|
||||
# Validate before committing so a malformed notification leaves
|
||||
# state untouched; then fall through so a registered handler
|
||||
# observes an initialized connection.
|
||||
if params is not None:
|
||||
try:
|
||||
NotificationParams.model_validate(params, by_name=False)
|
||||
except ValidationError:
|
||||
logger.warning("dropped %r: malformed params", method)
|
||||
return
|
||||
self.connection.initialized.set()
|
||||
elif not self.connection.initialize_accepted:
|
||||
logger.debug("dropped %s: received before initialization", method)
|
||||
return
|
||||
entry = self.server.get_notification_handler(method)
|
||||
if entry is None:
|
||||
logger.debug("no handler for notification %s", method)
|
||||
return
|
||||
# Same absent-params contract as requests.
|
||||
try:
|
||||
typed_params = entry.params_type.model_validate({} if params is None else params, by_name=False)
|
||||
except ValidationError:
|
||||
logger.warning("dropped %r: malformed params", method)
|
||||
return
|
||||
await entry.handler(ctx, typed_params)
|
||||
|
||||
call = self._compose_server_middleware(ctx, method, params, _inner)
|
||||
try:
|
||||
await call()
|
||||
except Exception:
|
||||
# A crashing handler must not cancel the dispatcher's task group;
|
||||
# middleware saw the raise out of call_next() first.
|
||||
logger.exception("notification handler for %r raised", method)
|
||||
|
||||
def _compose_server_middleware(
|
||||
self,
|
||||
ctx: ServerRequestContext[LifespanT, Any],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
inner: CallNext,
|
||||
) -> CallNext:
|
||||
"""Wrap `inner` in `Server.middleware`, outermost-first.
|
||||
|
||||
Shared by `_on_request` and `_on_notify` so the same middleware chain
|
||||
observes every inbound message.
|
||||
"""
|
||||
call = inner
|
||||
for mw in reversed(self.server.middleware):
|
||||
call = partial(mw, ctx, method, params, call)
|
||||
return call
|
||||
|
||||
def _make_context(
|
||||
self, dctx: DispatchContext[TransportContext], meta: RequestParamsMeta | None
|
||||
) -> ServerRequestContext[LifespanT, Any]:
|
||||
# TODO(maxisbey): remove for Context rework. Reads the SHTTP per-request
|
||||
# data off the raw `dctx.message_metadata` carrier; replace with the
|
||||
# per-transport context once that lands.
|
||||
md = dctx.message_metadata
|
||||
if isinstance(md, ServerMessageMetadata):
|
||||
request = md.request_context
|
||||
close_sse_stream = md.close_sse_stream
|
||||
close_standalone_sse_stream = md.close_standalone_sse_stream
|
||||
else:
|
||||
request = close_sse_stream = close_standalone_sse_stream = None
|
||||
return ServerRequestContext(
|
||||
session=self.session,
|
||||
lifespan_context=self.lifespan_state,
|
||||
request_id=dctx.request_id,
|
||||
meta=meta,
|
||||
request=request,
|
||||
close_sse_stream=close_sse_stream,
|
||||
close_standalone_sse_stream=close_standalone_sse_stream,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _negotiate_initialize(params: Mapping[str, Any] | None) -> tuple[InitializeRequestParams, str]:
|
||||
"""Validate `initialize` params and pick the protocol version."""
|
||||
init = InitializeRequestParams.model_validate(params or {}, by_name=False)
|
||||
requested = init.protocol_version
|
||||
negotiated = requested if requested in SUPPORTED_PROTOCOL_VERSIONS else LATEST_PROTOCOL_VERSION
|
||||
return init, negotiated
|
||||
|
||||
def _handle_initialize(self, params: Mapping[str, Any] | None) -> InitializeResult:
|
||||
"""Build the `initialize` result; state commits later in `_on_request`."""
|
||||
_, negotiated = self._negotiate_initialize(params)
|
||||
assert self.init_options is not None
|
||||
opts = self.init_options
|
||||
return InitializeResult(
|
||||
protocol_version=negotiated,
|
||||
capabilities=opts.capabilities,
|
||||
server_info=Implementation(
|
||||
name=opts.server_name,
|
||||
title=opts.title,
|
||||
description=opts.description,
|
||||
version=opts.server_version,
|
||||
website_url=opts.website_url,
|
||||
icons=opts.icons,
|
||||
),
|
||||
instructions=opts.instructions,
|
||||
)
|
||||
|
|
@ -1,188 +1,104 @@
|
|||
"""ServerSession Module
|
||||
"""`ServerSession`: server-to-client requests and notifications.
|
||||
|
||||
This module provides the ServerSession class, which manages communication between the
|
||||
server and client in the MCP (Model Context Protocol) framework. It is most commonly
|
||||
used in MCP servers to interact with the client.
|
||||
A thin proxy over `JSONRPCDispatcher` and `Connection`. One instance per
|
||||
client connection (built by `ServerRunner`). Handlers reach it as
|
||||
`ctx.session` and use the typed helpers (`create_message`, `elicit_form`,
|
||||
`send_log_message`, ...) to call back to the client.
|
||||
|
||||
Common usage pattern:
|
||||
```
|
||||
async def handle_call_tool(ctx: RequestContext, params: CallToolRequestParams) -> CallToolResult:
|
||||
# Check client capabilities before proceeding
|
||||
if ctx.session.check_client_capability(
|
||||
types.ClientCapabilities(experimental={"advanced_tools": dict()})
|
||||
):
|
||||
result = await perform_advanced_tool_operation(params.arguments)
|
||||
else:
|
||||
result = await perform_basic_tool_operation(params.arguments)
|
||||
return result
|
||||
|
||||
async def handle_list_prompts(ctx: RequestContext, params) -> ListPromptsResult:
|
||||
if ctx.session.client_params:
|
||||
return ListPromptsResult(prompts=generate_custom_prompts(ctx.session.client_params))
|
||||
return ListPromptsResult(prompts=default_prompts)
|
||||
|
||||
server = Server(name, on_call_tool=handle_call_tool, on_list_prompts=handle_list_prompts)
|
||||
```
|
||||
|
||||
The ServerSession class is typically used internally by the Server class and should not
|
||||
be instantiated directly by users of the MCP framework.
|
||||
The receive-loop, initialize handling, and per-request task isolation that
|
||||
used to live here are now owned by `JSONRPCDispatcher` and `ServerRunner`.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, TypeVar, overload
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream
|
||||
from pydantic import AnyUrl, TypeAdapter
|
||||
from pydantic import AnyUrl, BaseModel
|
||||
|
||||
from mcp import types
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
|
||||
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
||||
from mcp.shared.exceptions import StatelessModeNotSupported
|
||||
from mcp.shared.message import ServerMessageMetadata, SessionMessage
|
||||
from mcp.shared.session import (
|
||||
BaseSession,
|
||||
RequestResponder,
|
||||
)
|
||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||
from mcp.shared.dispatcher import CallOptions, ProgressFnT
|
||||
from mcp.shared.exceptions import NoBackChannelError, StatelessModeNotSupported
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.message import ServerMessageMetadata
|
||||
|
||||
__all__ = ["ServerSession"]
|
||||
|
||||
ResultT = TypeVar("ResultT", bound=BaseModel)
|
||||
|
||||
|
||||
class InitializationState(Enum):
|
||||
NotInitialized = 1
|
||||
Initializing = 2
|
||||
Initialized = 3
|
||||
class ServerSession:
|
||||
"""Connection-scoped proxy for server-to-client requests and notifications.
|
||||
|
||||
|
||||
ServerSessionT = TypeVar("ServerSessionT", bound="ServerSession")
|
||||
|
||||
ServerRequestResponder = (
|
||||
RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception
|
||||
)
|
||||
|
||||
|
||||
class ServerSession(
|
||||
BaseSession[
|
||||
types.ServerRequest,
|
||||
types.ServerNotification,
|
||||
types.ServerResult,
|
||||
types.ClientRequest,
|
||||
types.ClientNotification,
|
||||
]
|
||||
):
|
||||
_initialized: InitializationState = InitializationState.NotInitialized
|
||||
_client_params: types.InitializeRequestParams | None = None
|
||||
`send_request` / `send_notification` model-dump their argument and forward
|
||||
to the dispatcher; the typed helpers below are unchanged from the previous
|
||||
implementation and only call those two methods.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
read_stream: ReadStream[SessionMessage | Exception],
|
||||
write_stream: WriteStream[SessionMessage],
|
||||
init_options: InitializationOptions,
|
||||
dispatcher: JSONRPCDispatcher[Any],
|
||||
connection: Connection,
|
||||
*,
|
||||
stateless: bool = False,
|
||||
) -> None:
|
||||
super().__init__(read_stream, write_stream)
|
||||
self._dispatcher = dispatcher
|
||||
self._connection = connection
|
||||
self._stateless = stateless
|
||||
self._initialization_state = (
|
||||
InitializationState.Initialized if stateless else InitializationState.NotInitialized
|
||||
)
|
||||
|
||||
self._init_options = init_options
|
||||
self._incoming_message_stream_writer, self._incoming_message_stream_reader = anyio.create_memory_object_stream[
|
||||
ServerRequestResponder
|
||||
](0)
|
||||
self._exit_stack.push_async_callback(lambda: self._incoming_message_stream_reader.aclose())
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> TypeAdapter[types.ClientRequest]:
|
||||
return types.client_request_adapter
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> TypeAdapter[types.ClientNotification]:
|
||||
return types.client_notification_adapter
|
||||
|
||||
@property
|
||||
def client_params(self) -> types.InitializeRequestParams | None:
|
||||
return self._client_params
|
||||
"""The client's `initialize` request params; `None` before initialization."""
|
||||
return self._connection.client_params
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
request: types.ServerRequest,
|
||||
result_type: type[ResultT],
|
||||
request_read_timeout_seconds: float | None = None,
|
||||
metadata: ServerMessageMetadata | None = None,
|
||||
progress_callback: ProgressFnT | None = None,
|
||||
) -> ResultT:
|
||||
"""Send a typed server-to-client request and validate the result.
|
||||
|
||||
`metadata.related_request_id` (when supplied) routes the outgoing
|
||||
message onto the originating request's response stream over
|
||||
streamable HTTP; it is the only metadata field honored here.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: If there is no related request to ride on and
|
||||
the connection has no standalone channel (stateless HTTP), so
|
||||
a response could never arrive.
|
||||
pydantic.ValidationError: The peer's result does not match `result_type`.
|
||||
"""
|
||||
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
opts: CallOptions = {}
|
||||
if request_read_timeout_seconds:
|
||||
opts["timeout"] = request_read_timeout_seconds
|
||||
if progress_callback is not None:
|
||||
opts["on_progress"] = progress_callback
|
||||
related = metadata.related_request_id if metadata is not None else None
|
||||
if related is None and not self._connection.has_standalone_channel:
|
||||
# Fail fast instead of parking forever on a response that cannot
|
||||
# arrive; matches `Connection.send_raw_request`.
|
||||
raise NoBackChannelError(data["method"])
|
||||
result = await self._dispatcher.send_raw_request(
|
||||
data["method"], data.get("params"), opts or None, _related_request_id=related
|
||||
)
|
||||
return result_type.model_validate(result, by_name=False)
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
notification: types.ServerNotification,
|
||||
related_request_id: types.RequestId | None = None,
|
||||
) -> None:
|
||||
"""Send a typed server-to-client notification."""
|
||||
data = notification.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
await self._dispatcher.notify(data["method"], data.get("params"), _related_request_id=related_request_id)
|
||||
|
||||
def check_client_capability(self, capability: types.ClientCapabilities) -> bool:
|
||||
"""Check if the client supports a specific capability."""
|
||||
if self._client_params is None: # pragma: lax no cover
|
||||
return False
|
||||
|
||||
client_caps = self._client_params.capabilities
|
||||
|
||||
if capability.roots is not None: # pragma: lax no cover
|
||||
if client_caps.roots is None:
|
||||
return False
|
||||
if capability.roots.list_changed and not client_caps.roots.list_changed:
|
||||
return False
|
||||
|
||||
if capability.sampling is not None: # pragma: lax no cover
|
||||
if client_caps.sampling is None:
|
||||
return False
|
||||
if capability.sampling.context is not None and client_caps.sampling.context is None:
|
||||
return False
|
||||
if capability.sampling.tools is not None and client_caps.sampling.tools is None:
|
||||
return False
|
||||
|
||||
if capability.elicitation is not None and client_caps.elicitation is None: # pragma: lax no cover
|
||||
return False
|
||||
|
||||
if capability.experimental is not None: # pragma: lax no cover
|
||||
if client_caps.experimental is None:
|
||||
return False
|
||||
for exp_key, exp_value in capability.experimental.items():
|
||||
if exp_key not in client_caps.experimental or client_caps.experimental[exp_key] != exp_value:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def _receive_loop(self) -> None:
|
||||
async with self._incoming_message_stream_writer:
|
||||
await super()._receive_loop()
|
||||
|
||||
async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]):
|
||||
match responder.request:
|
||||
case types.InitializeRequest(params=params):
|
||||
requested_version = params.protocol_version
|
||||
self._initialization_state = InitializationState.Initializing
|
||||
self._client_params = params
|
||||
with responder:
|
||||
await responder.respond(
|
||||
types.InitializeResult(
|
||||
protocol_version=requested_version
|
||||
if requested_version in SUPPORTED_PROTOCOL_VERSIONS
|
||||
else types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=self._init_options.capabilities,
|
||||
server_info=types.Implementation(
|
||||
name=self._init_options.server_name,
|
||||
title=self._init_options.title,
|
||||
description=self._init_options.description,
|
||||
version=self._init_options.server_version,
|
||||
website_url=self._init_options.website_url,
|
||||
icons=self._init_options.icons,
|
||||
),
|
||||
instructions=self._init_options.instructions,
|
||||
)
|
||||
)
|
||||
self._initialization_state = InitializationState.Initialized
|
||||
case types.PingRequest():
|
||||
# Ping requests are allowed at any time
|
||||
pass
|
||||
case _:
|
||||
if self._initialization_state != InitializationState.Initialized:
|
||||
raise RuntimeError("Received request before initialization was complete")
|
||||
|
||||
async def _received_notification(self, notification: types.ClientNotification) -> None:
|
||||
# Need this to avoid ASYNC910
|
||||
await anyio.lowlevel.checkpoint()
|
||||
match notification:
|
||||
case types.InitializedNotification():
|
||||
self._initialization_state = InitializationState.Initialized
|
||||
case _:
|
||||
if self._initialization_state != InitializationState.Initialized: # pragma: no cover
|
||||
raise RuntimeError("Received notification before initialization was complete")
|
||||
return self._connection.check_capability(capability)
|
||||
|
||||
async def send_log_message(
|
||||
self,
|
||||
|
|
@ -293,7 +209,7 @@ class ServerSession(
|
|||
"""
|
||||
if self._stateless:
|
||||
raise StatelessModeNotSupported(method="sampling")
|
||||
client_caps = self._client_params.capabilities if self._client_params else None
|
||||
client_caps = self.client_params.capabilities if self.client_params else None
|
||||
validate_sampling_tools(client_caps, tools, tool_choice)
|
||||
validate_tool_use_result_messages(messages)
|
||||
|
||||
|
|
@ -313,7 +229,6 @@ class ServerSession(
|
|||
)
|
||||
metadata_obj = ServerMessageMetadata(related_request_id=related_request_id)
|
||||
|
||||
# Use different result types based on whether tools are provided
|
||||
if tools is not None:
|
||||
return await self.send_request(
|
||||
request=request,
|
||||
|
|
@ -488,10 +403,3 @@ class ServerSession(
|
|||
),
|
||||
related_request_id,
|
||||
)
|
||||
|
||||
async def _handle_incoming(self, req: ServerRequestResponder) -> None:
|
||||
await self._incoming_message_stream_writer.send(req)
|
||||
|
||||
@property
|
||||
def incoming_messages(self) -> MemoryObjectReceiveStream[ServerRequestResponder]:
|
||||
return self._incoming_message_stream_reader
|
||||
|
|
|
|||
|
|
@ -633,7 +633,10 @@ class StreamableHTTPServerTransport:
|
|||
finally:
|
||||
await sse_stream_reader.aclose()
|
||||
|
||||
except Exception as err:
|
||||
except Exception as err: # pragma: lax no cover
|
||||
# Reached only when something raises during POST handling outside
|
||||
# the per-SSE-stream guard above; whether tests reach this depends
|
||||
# on client teardown timing.
|
||||
logger.exception("Error handling POST request")
|
||||
response = self._create_error_response(
|
||||
f"Error handling POST request: {err}",
|
||||
|
|
@ -641,7 +644,7 @@ class StreamableHTTPServerTransport:
|
|||
INTERNAL_ERROR,
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
if writer: # pragma: no cover
|
||||
if writer:
|
||||
await writer.send(Exception(err))
|
||||
return # pragma: no cover
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from mcp.server.streamable_http import (
|
|||
StreamableHTTPServerTransport,
|
||||
)
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.types import INVALID_REQUEST, ErrorData, JSONRPCError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -139,6 +140,7 @@ class StreamableHTTPSessionManager:
|
|||
# Clear any remaining server instances
|
||||
self._server_instances.clear()
|
||||
self._session_owners.clear()
|
||||
await resync_tracer()
|
||||
|
||||
async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
"""Process ASGI request with proper session handling and transport setup.
|
||||
|
|
|
|||
19
src/mcp/shared/_compat.py
Normal file
19
src/mcp/shared/_compat.py
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
"""Workarounds for CPython interpreter bugs the SDK papers over."""
|
||||
|
||||
import anyio.lowlevel
|
||||
|
||||
__all__ = ["resync_tracer"]
|
||||
|
||||
|
||||
async def resync_tracer() -> None:
|
||||
"""Resync coverage tracing after a cancelled task-group join.
|
||||
|
||||
A cancel delivered at a join resumes the awaiting coroutine chain via
|
||||
`coro.throw()`; on CPython 3.11 (python/cpython#106749) that drops the
|
||||
`'call'` trace events for the outer frames and desyncs coverage's CTracer
|
||||
until the chain next suspends and resumes normally. Yielding once here
|
||||
resumes via `.send()`, which re-stamps the missing events. Shielded so a
|
||||
pending outer cancel is not re-delivered at this point; behaviorally a
|
||||
no-op. Delete this module when Python 3.11 support ends (EOL 2027-10).
|
||||
"""
|
||||
await anyio.lowlevel.cancel_shielded_checkpoint()
|
||||
|
|
@ -20,9 +20,18 @@ def otel_span(
|
|||
kind: SpanKind,
|
||||
attributes: dict[str, Any] | None = None,
|
||||
context: Context | None = None,
|
||||
record_exception: bool = True,
|
||||
set_status_on_exception: bool = True,
|
||||
) -> Iterator[Any]:
|
||||
"""Create an OTel span."""
|
||||
with _tracer.start_as_current_span(name, kind=kind, attributes=attributes, context=context) as span:
|
||||
with _tracer.start_as_current_span(
|
||||
name,
|
||||
kind=kind,
|
||||
attributes=attributes,
|
||||
context=context,
|
||||
record_exception=record_exception,
|
||||
set_status_on_exception=set_status_on_exception,
|
||||
) as span:
|
||||
yield span
|
||||
|
||||
|
||||
|
|
@ -31,6 +40,13 @@ def inject_trace_context(meta: dict[str, Any]) -> None:
|
|||
inject(meta)
|
||||
|
||||
|
||||
def extract_trace_context(meta: dict[str, Any]) -> Context:
|
||||
"""Extract W3C trace context from a `_meta` dict."""
|
||||
return extract(meta)
|
||||
def extract_trace_context(meta: dict[str, Any]) -> Context | None:
|
||||
"""Extract W3C trace context from a `_meta` dict.
|
||||
|
||||
Returns `None` when the carrier is malformed; telemetry parsing must
|
||||
never fail the request it annotates.
|
||||
"""
|
||||
try:
|
||||
return extract(meta)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
|
|
|||
85
src/mcp/shared/context.py
Normal file
85
src/mcp/shared/context.py
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
"""`BaseContext` - the user-facing per-request context.
|
||||
|
||||
Composition over a `DispatchContext`: forwards the transport metadata, the
|
||||
back-channel (`send_raw_request`/`notify`), progress reporting, and the cancel
|
||||
event. Adds `meta` (the inbound request's `_meta` field).
|
||||
|
||||
Satisfies `Outbound`, so `ClientPeer` can wrap it. Shared between client and
|
||||
server: the server's `Context` extends this with `lifespan`/`connection`;
|
||||
`ClientContext` is just an alias.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Generic
|
||||
|
||||
import anyio
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.shared.dispatcher import CallOptions, DispatchContext
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import RequestParamsMeta
|
||||
|
||||
__all__ = ["BaseContext"]
|
||||
|
||||
TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext, covariant=True)
|
||||
|
||||
|
||||
class BaseContext(Generic[TransportT]):
|
||||
"""Per-request context wrapping a `DispatchContext`.
|
||||
|
||||
`ServerRunner` constructs one per inbound request and passes it to the
|
||||
user's handler.
|
||||
"""
|
||||
|
||||
def __init__(self, dctx: DispatchContext[TransportT], meta: RequestParamsMeta | None = None) -> None:
|
||||
self._dctx = dctx
|
||||
self._meta = meta
|
||||
|
||||
@property
|
||||
def transport(self) -> TransportT:
|
||||
"""Transport-specific metadata for this inbound request."""
|
||||
return self._dctx.transport
|
||||
|
||||
@property
|
||||
def cancel_requested(self) -> anyio.Event:
|
||||
"""Set when the peer sends `notifications/cancelled` for this request."""
|
||||
return self._dctx.cancel_requested
|
||||
|
||||
@property
|
||||
def can_send_request(self) -> bool:
|
||||
"""Whether the back-channel can currently deliver server-initiated requests.
|
||||
|
||||
`False` when the transport has no back-channel, or when the underlying
|
||||
dispatch context has been closed because the inbound request finished.
|
||||
"""
|
||||
return self._dctx.can_send_request
|
||||
|
||||
@property
|
||||
def meta(self) -> RequestParamsMeta | None:
|
||||
"""The inbound request's `_meta` field, if present."""
|
||||
return self._meta
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Send a request to the peer on the back-channel.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: `can_send_request` is `False`.
|
||||
"""
|
||||
return await self._dctx.send_raw_request(method, params, opts)
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
"""Send a notification to the peer on the back-channel."""
|
||||
await self._dctx.notify(method, params)
|
||||
|
||||
async def report_progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
||||
"""Report progress for this request, if the peer supplied a progress token.
|
||||
|
||||
A no-op when no token was supplied.
|
||||
"""
|
||||
await self._dctx.progress(progress, total, message)
|
||||
197
src/mcp/shared/direct_dispatcher.py
Normal file
197
src/mcp/shared/direct_dispatcher.py
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
"""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
|
||||
|
||||
Unlike `JSONRPCDispatcher`, exceptions raised in a handler propagate directly
|
||||
to the caller - there is no exception-to-`ErrorData` boundary here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import anyio.abc
|
||||
from pydantic import ValidationError
|
||||
|
||||
from mcp.shared.dispatcher import CallOptions, OnNotify, OnRequest, ProgressFnT
|
||||
from mcp.shared.exceptions import MCPError, NoBackChannelError
|
||||
from mcp.shared.message import MessageMetadata
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import INTERNAL_ERROR, INVALID_PARAMS, REQUEST_TIMEOUT, RequestId
|
||||
|
||||
__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
|
||||
"""Always `None`: direct dispatch has no wire-level request id."""
|
||||
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) -> 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.
|
||||
"""
|
||||
|
||||
def __init__(self, transport_ctx: TransportContext):
|
||||
self._transport_ctx = transport_ctx
|
||||
self._peer: DirectDispatcher | None = None
|
||||
self._on_request: OnRequest | None = None
|
||||
self._on_notify: OnNotify | None = None
|
||||
self._ready = anyio.Event()
|
||||
self._closed = anyio.Event()
|
||||
|
||||
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]:
|
||||
if self._peer is None:
|
||||
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
|
||||
return await self._peer._dispatch_request(method, params, opts)
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
if self._peer is 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,
|
||||
*,
|
||||
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:
|
||||
self._closed.set()
|
||||
|
||||
def _make_context(self, on_progress: ProgressFnT | 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),
|
||||
_on_progress=on_progress,
|
||||
)
|
||||
|
||||
async def _dispatch_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None,
|
||||
) -> dict[str, Any]:
|
||||
await self._ready.wait()
|
||||
assert self._on_request is not None
|
||||
opts = opts or {}
|
||||
dctx = self._make_context(on_progress=opts.get("on_progress"))
|
||||
try:
|
||||
with anyio.fail_after(opts.get("timeout")):
|
||||
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:
|
||||
raise MCPError(code=INTERNAL_ERROR, message=str(e)) from e
|
||||
except TimeoutError:
|
||||
raise MCPError(
|
||||
code=REQUEST_TIMEOUT,
|
||||
message=f"Timed out after {opts.get('timeout')}s waiting for {method!r}",
|
||||
) from None
|
||||
|
||||
async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
await self._ready.wait()
|
||||
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,
|
||||
) -> 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.
|
||||
|
||||
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)
|
||||
server = DirectDispatcher(ctx)
|
||||
client.connect_to(server)
|
||||
server.connect_to(client)
|
||||
return client, server
|
||||
209
src/mcp/shared/dispatcher.py
Normal file
209
src/mcp/shared/dispatcher.py
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
"""Dispatcher Protocol - the call/return boundary between transports and handlers.
|
||||
|
||||
A Dispatcher turns a duplex message channel into two things:
|
||||
|
||||
* an outbound API: `send_raw_request(method, params)` and `notify(method, params)`
|
||||
* an inbound pump: `run(on_request, on_notify)` that drives the receive loop
|
||||
and invokes the supplied handlers for each incoming request/notification
|
||||
|
||||
It is deliberately *not* MCP-aware. Method names are strings, params and
|
||||
results are `dict[str, Any]`. The MCP type layer (request/result models,
|
||||
capability negotiation, `Context`) sits above this; the wire encoding
|
||||
(JSON-RPC, gRPC, in-process direct calls) sits below it.
|
||||
|
||||
See `JSONRPCDispatcher` for the production implementation and
|
||||
`DirectDispatcher` for an in-memory implementation used in tests and for
|
||||
embedding a server in-process.
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable
|
||||
|
||||
import anyio
|
||||
import anyio.abc
|
||||
|
||||
from mcp.shared.message import MessageMetadata
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import RequestId
|
||||
|
||||
__all__ = [
|
||||
"CallOptions",
|
||||
"DispatchContext",
|
||||
"DispatchMiddleware",
|
||||
"Dispatcher",
|
||||
"OnNotify",
|
||||
"OnRequest",
|
||||
"Outbound",
|
||||
"ProgressFnT",
|
||||
]
|
||||
|
||||
TransportT_co = TypeVar("TransportT_co", bound=TransportContext, covariant=True)
|
||||
|
||||
|
||||
class ProgressFnT(Protocol):
|
||||
"""Callback invoked when a progress notification arrives for a pending request."""
|
||||
|
||||
async def __call__(self, progress: float, total: float | None, message: str | None) -> None: ...
|
||||
|
||||
|
||||
class CallOptions(TypedDict, total=False):
|
||||
"""Per-call options for `Outbound.send_raw_request`.
|
||||
|
||||
All keys are optional. Dispatchers ignore keys they do not understand.
|
||||
"""
|
||||
|
||||
timeout: float
|
||||
"""Seconds to wait for a result before raising and sending `notifications/cancelled`."""
|
||||
|
||||
on_progress: ProgressFnT
|
||||
"""Receive `notifications/progress` updates for this request."""
|
||||
|
||||
resumption_token: str
|
||||
"""Opaque token to resume a previously interrupted request.
|
||||
|
||||
Client-side, streamable-HTTP only. Ignored by server dispatchers and other
|
||||
transports, and also ignored (with a debug log) for requests sent from a
|
||||
`DispatchContext`, where routing onto the inbound request's stream takes
|
||||
precedence. Supports protocol version 2025-11-25 and earlier; SSE-stream
|
||||
resumption is removed in the next protocol revision.
|
||||
"""
|
||||
|
||||
on_resumption_token: Callable[[str], Awaitable[None]]
|
||||
"""Receive a resumption token when the transport issues one for this request.
|
||||
|
||||
Client-side, streamable-HTTP only. Ignored by server dispatchers and other
|
||||
transports, and also ignored (with a debug log) for requests sent from a
|
||||
`DispatchContext`, where routing onto the inbound request's stream takes
|
||||
precedence. Supports protocol version 2025-11-25 and earlier; SSE-stream
|
||||
resumption is removed in the next protocol revision.
|
||||
"""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Outbound(Protocol):
|
||||
"""Anything that can send requests and notifications to the peer.
|
||||
|
||||
Both `Dispatcher` (top-level outbound) and `DispatchContext` (back-channel
|
||||
during an inbound request) extend this. The MCP type layer (`ClientPeer`,
|
||||
`Connection`) builds typed `send_request` / convenience methods on top of
|
||||
this raw channel.
|
||||
"""
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Send a request and await its raw result dict.
|
||||
|
||||
`opts` carries per-call `timeout` / `on_progress` / resumption hints;
|
||||
see `CallOptions`.
|
||||
|
||||
Raises:
|
||||
MCPError: If the peer responded with an error, or the handler
|
||||
raised. Implementations normalize all handler exceptions to
|
||||
`MCPError` so callers see a single exception type.
|
||||
"""
|
||||
...
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
"""Send a fire-and-forget notification."""
|
||||
...
|
||||
|
||||
|
||||
class DispatchContext(Outbound, Protocol[TransportT_co]):
|
||||
"""Per-request context handed to `on_request` / `on_notify`.
|
||||
|
||||
Carries the transport metadata for the inbound message and provides the
|
||||
back-channel for sending requests/notifications to the peer while handling
|
||||
it. `send_raw_request` raises `NoBackChannelError` if `can_send_request`
|
||||
is `False`.
|
||||
"""
|
||||
|
||||
@property
|
||||
def transport(self) -> TransportT_co:
|
||||
"""Transport-specific metadata for this inbound message."""
|
||||
...
|
||||
|
||||
@property
|
||||
def can_send_request(self) -> bool:
|
||||
"""Whether the back-channel can currently deliver server-initiated requests.
|
||||
|
||||
`False` when the transport has no back-channel, or when this context has
|
||||
been closed (the inbound request finished). `send_raw_request` raises
|
||||
`NoBackChannelError` exactly when this is `False`.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
def request_id(self) -> RequestId | None:
|
||||
"""The id of the inbound request, or `None` for a notification.
|
||||
|
||||
For JSON-RPC this is the wire `id` field. Handlers thread it through
|
||||
as `related_request_id` on outbound notifications so HTTP transports
|
||||
can route them onto the originating request's response stream.
|
||||
"""
|
||||
...
|
||||
|
||||
@property
|
||||
def message_metadata(self) -> MessageMetadata:
|
||||
"""The metadata the transport attached to this inbound message, if any.
|
||||
|
||||
This is `SessionMessage.metadata` passed through verbatim: HTTP
|
||||
transports attach `ServerMessageMetadata` (the HTTP request, SSE
|
||||
stream-close callbacks); stdio and in-memory dispatch attach nothing.
|
||||
Tied to the `SessionMessage` wire format - goes away when transports
|
||||
stop delivering messages that way.
|
||||
"""
|
||||
# TODO(maxisbey): remove for context rework
|
||||
...
|
||||
|
||||
@property
|
||||
def cancel_requested(self) -> anyio.Event:
|
||||
"""Set when the peer sends `notifications/cancelled` for this request."""
|
||||
...
|
||||
|
||||
async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
||||
"""Report progress for the inbound request, if the peer supplied a progress token.
|
||||
|
||||
A no-op when no token was supplied.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
OnRequest = Callable[[DispatchContext[TransportContext], str, Mapping[str, Any] | None], Awaitable[dict[str, Any]]]
|
||||
"""Handler for inbound requests: `(ctx, method, params) -> result`. Raise `MCPError` to send an error response."""
|
||||
|
||||
OnNotify = Callable[[DispatchContext[TransportContext], str, Mapping[str, Any] | None], Awaitable[None]]
|
||||
"""Handler for inbound notifications: `(ctx, method, params)`."""
|
||||
|
||||
DispatchMiddleware = Callable[[OnRequest], OnRequest]
|
||||
"""Wraps an `OnRequest` to produce another `OnRequest`. Applied outermost-first."""
|
||||
|
||||
|
||||
class Dispatcher(Outbound, Protocol[TransportT_co]):
|
||||
"""A duplex request/notification channel with call-return semantics.
|
||||
|
||||
Implementations own correlation of outbound requests to inbound results, the
|
||||
receive loop, per-request concurrency, and cancellation/progress wiring.
|
||||
"""
|
||||
|
||||
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)`.
|
||||
"""
|
||||
...
|
||||
|
|
@ -2,7 +2,7 @@ from __future__ import annotations
|
|||
|
||||
from typing import Any, cast
|
||||
|
||||
from mcp.types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError
|
||||
from mcp.types import INVALID_REQUEST, URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData, JSONRPCError
|
||||
|
||||
|
||||
class MCPError(Exception):
|
||||
|
|
@ -41,6 +41,25 @@ class MCPError(Exception):
|
|||
return self.message
|
||||
|
||||
|
||||
class NoBackChannelError(MCPError):
|
||||
"""Raised when sending a server-initiated request over a transport that cannot deliver it.
|
||||
|
||||
Stateless HTTP and JSON-response-mode HTTP have no channel for the server to
|
||||
push requests (sampling, elicitation, roots/list) to the client. This is
|
||||
raised by `DispatchContext.send_raw_request` when `can_send_request` is
|
||||
`False`, and serializes to an `INVALID_REQUEST` error response.
|
||||
"""
|
||||
|
||||
def __init__(self, method: str):
|
||||
super().__init__(
|
||||
code=INVALID_REQUEST,
|
||||
message=(
|
||||
f"Cannot send {method!r}: this transport context has no back-channel for server-initiated requests."
|
||||
),
|
||||
)
|
||||
self.method = method
|
||||
|
||||
|
||||
class StatelessModeNotSupported(RuntimeError):
|
||||
"""Raised when attempting to use a method that is not supported in stateless mode.
|
||||
|
||||
|
|
|
|||
742
src/mcp/shared/jsonrpc_dispatcher.py
Normal file
742
src/mcp/shared/jsonrpc_dispatcher.py
Normal file
|
|
@ -0,0 +1,742 @@
|
|||
"""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
|
||||
|
|
@ -5,6 +5,7 @@ from __future__ import annotations
|
|||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
|
||||
from mcp.shared.message import SessionMessage
|
||||
|
||||
|
|
@ -28,3 +29,5 @@ async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageS
|
|||
|
||||
async with server_to_client_receive, client_to_server_send, client_to_server_receive, server_to_client_send:
|
||||
yield client_streams, server_streams
|
||||
# Heals caller-driven cancels; closing memory streams never suspends.
|
||||
await resync_tracer()
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue