Server-side 2026-07-28 stateless support: classifier, driver split, server/discover (#2928)
This commit is contained in:
parent
44724284b3
commit
2397319a68
1 changed files with 2256 additions and 1088 deletions
|
|
@ -1,14 +1,12 @@
|
|||
"""`ServerRunner` - per-connection orchestrator over a `Dispatcher`.
|
||||
"""`ServerRunner` - the per-connection handler kernel.
|
||||
|
||||
`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` bridges the dispatch layer (`on_request` / `on_notify`, untyped
|
||||
dicts) and the user's handler layer (typed `Context`, typed params). It is a
|
||||
pure kernel: it holds a pre-populated `Connection` and reads
|
||||
`connection.protocol_version` / `connection.outbound` as facts. Driving a
|
||||
dispatcher loop and tearing down the connection live in the free-function
|
||||
drivers (`serve_connection`, `serve_loop`, `serve_one`); the entry constructs
|
||||
the `Connection`, the driver tears it down.
|
||||
|
||||
`ServerRunner` holds a `Server` directly - `Server` is the registry.
|
||||
"""
|
||||
|
|
@ -16,11 +14,12 @@ typed params). One instance per client connection. It:
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial, reduce
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from dataclasses import KW_ONLY, dataclass
|
||||
from functools import cached_property, partial, reduce
|
||||
from typing import TYPE_CHECKING, Any, Generic, cast
|
||||
|
||||
import anyio
|
||||
import anyio.abc
|
||||
from opentelemetry.trace import SpanKind, StatusCode
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
|
@ -31,9 +30,11 @@ from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, Server
|
|||
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, Dispatcher, DispatchMiddleware, OnRequest
|
||||
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
||||
from mcp.shared.dispatcher import DispatchContext, Dispatcher, DispatchMiddleware, OnNotify, OnRequest
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.message import MessageMetadata, ServerMessageMetadata
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher, handler_exception_to_error_data
|
||||
from mcp.shared.message import ServerMessageMetadata, SessionMessage
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||
from mcp.types import (
|
||||
|
|
@ -41,11 +42,14 @@ from mcp.types import (
|
|||
INVALID_PARAMS,
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
METHOD_NOT_FOUND,
|
||||
PROTOCOL_VERSION_META_KEY,
|
||||
ErrorData,
|
||||
Implementation,
|
||||
InitializeRequestParams,
|
||||
InitializeResult,
|
||||
JSONRPCError,
|
||||
JSONRPCRequest,
|
||||
JSONRPCResponse,
|
||||
RequestId,
|
||||
RequestParams,
|
||||
RequestParamsMeta,
|
||||
)
|
||||
|
|
@ -54,7 +58,17 @@ from mcp.types import methods as _methods
|
|||
if TYPE_CHECKING:
|
||||
from mcp.server.lowlevel.server import Server
|
||||
|
||||
__all__ = ["CallNext", "ServerMiddleware", "ServerRunner", "otel_middleware"]
|
||||
__all__ = [
|
||||
"CallNext",
|
||||
"ServerMiddleware",
|
||||
"ServerRunner",
|
||||
"aclose_shielded",
|
||||
"otel_middleware",
|
||||
"serve_connection",
|
||||
"serve_loop",
|
||||
"serve_one",
|
||||
"to_jsonrpc_response",
|
||||
]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -64,8 +78,8 @@ 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."""
|
||||
"""Bound for `aclose_shielded`'s exit-stack unwind; a hung cleanup callback
|
||||
must not wedge shutdown."""
|
||||
|
||||
|
||||
def _extract_meta(params: Mapping[str, Any] | None) -> RequestParamsMeta | None:
|
||||
|
|
@ -79,30 +93,6 @@ def _extract_meta(params: Mapping[str, Any] | None) -> RequestParamsMeta | None:
|
|||
return None
|
||||
|
||||
|
||||
def _resolve_protocol_version(
|
||||
negotiated: str | None,
|
||||
meta: RequestParamsMeta | None,
|
||||
md: MessageMetadata,
|
||||
) -> str:
|
||||
"""Resolve the protocol version for this inbound message.
|
||||
|
||||
Handshake-committed value wins; else per-request `_meta`, else the
|
||||
transport hint. Unsupported values fall through so surface validation
|
||||
never sees them.
|
||||
"""
|
||||
if negotiated is not None:
|
||||
return negotiated
|
||||
if meta is not None:
|
||||
v = meta.get(PROTOCOL_VERSION_META_KEY)
|
||||
if isinstance(v, str) and v in SUPPORTED_PROTOCOL_VERSIONS:
|
||||
return v
|
||||
if isinstance(md, ServerMessageMetadata):
|
||||
hint = md.protocol_version
|
||||
if hint is not None and hint in SUPPORTED_PROTOCOL_VERSIONS:
|
||||
return hint
|
||||
return "2025-11-25"
|
||||
|
||||
|
||||
def otel_middleware(next_on_request: OnRequest) -> OnRequest:
|
||||
"""Dispatch-tier middleware that wraps each request in an OpenTelemetry span.
|
||||
|
||||
|
|
@ -169,73 +159,73 @@ def _dump_result(result: Any) -> dict[str, Any]:
|
|||
raise TypeError(f"handler returned {type(result).__name__}; expected BaseModel, dict, or None")
|
||||
|
||||
|
||||
async def aclose_shielded(connection: Connection) -> None:
|
||||
"""Unwind ``connection.exit_stack`` under a shielded, bounded scope.
|
||||
|
||||
Called from a driver's ``finally``: the shield lets per-connection cleanup
|
||||
callbacks run even when the driver itself is being cancelled, the
|
||||
`_EXIT_STACK_CLOSE_TIMEOUT` bound stops a hung callback wedging shutdown,
|
||||
and a raising callback is logged-and-swallowed so it never masks the
|
||||
driver's own exception.
|
||||
"""
|
||||
with anyio.move_on_after(_EXIT_STACK_CLOSE_TIMEOUT, shield=True) as scope:
|
||||
try:
|
||||
await connection.exit_stack.aclose()
|
||||
except Exception:
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
async def to_jsonrpc_response(request_id: RequestId, coro: Awaitable[dict[str, Any]]) -> JSONRPCResponse | JSONRPCError:
|
||||
"""Await ``coro`` and wrap its outcome as the JSON-RPC reply for ``request_id``.
|
||||
|
||||
The exception-to-wire boundary for the request-per-call drivers
|
||||
(`serve_one`, the modern HTTP entry). `MCPError` and `ValidationError`
|
||||
map via the shared `handler_exception_to_error_data` ladder; any other
|
||||
exception is logged and surfaced as `INTERNAL_ERROR` so handler internals
|
||||
never reach the wire.
|
||||
"""
|
||||
try:
|
||||
result = await coro
|
||||
except Exception as exc:
|
||||
error = handler_exception_to_error_data(exc)
|
||||
if error is None:
|
||||
logger.exception("request handler raised")
|
||||
error = ErrorData(code=INTERNAL_ERROR, message="Internal server error")
|
||||
return JSONRPCError(jsonrpc="2.0", id=request_id, error=error)
|
||||
return JSONRPCResponse(jsonrpc="2.0", id=request_id, result=result)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerRunner(Generic[LifespanT]):
|
||||
"""Per-connection orchestrator. One instance per client connection."""
|
||||
"""Per-connection handler kernel. One instance per client connection."""
|
||||
|
||||
server: Server[LifespanT]
|
||||
dispatcher: Dispatcher[Any]
|
||||
connection: Connection
|
||||
lifespan_state: LifespanT
|
||||
has_standalone_channel: bool
|
||||
_: KW_ONLY
|
||||
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])
|
||||
dispatch_middleware: Sequence[DispatchMiddleware] = (otel_middleware,)
|
||||
|
||||
connection: Connection = field(init=False)
|
||||
session: ServerSession = field(init=False)
|
||||
"""Connection-scoped: the same instance reaches every request as `ctx.session`."""
|
||||
@cached_property
|
||||
def on_request(self) -> OnRequest:
|
||||
"""`_on_request` wrapped in `dispatch_middleware`, outermost-first.
|
||||
|
||||
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()`.
|
||||
Dispatch-tier middleware sees raw `(dctx, method, params) -> dict` and
|
||||
wraps everything - initialize, METHOD_NOT_FOUND, validation failures
|
||||
included.
|
||||
"""
|
||||
return reduce(lambda h, mw: mw(h), reversed(self.dispatch_middleware), self._on_request)
|
||||
|
||||
@cached_property
|
||||
def on_notify(self) -> OnNotify:
|
||||
return self._on_notify
|
||||
|
||||
async def _on_request(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
|
|
@ -243,7 +233,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
params: Mapping[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
meta = _extract_meta(params)
|
||||
version = _resolve_protocol_version(self.connection.protocol_version, meta, dctx.message_metadata)
|
||||
version = self.connection.protocol_version
|
||||
ctx = self._make_context(dctx, meta, version)
|
||||
is_spec_method = method in _methods.SPEC_CLIENT_METHODS
|
||||
|
||||
|
|
@ -257,7 +247,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
_methods.validate_client_request(method, version, params)
|
||||
except KeyError:
|
||||
raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method) from None
|
||||
# TODO(maxisbey): the 2026-07-28 spec drops the handshake; this branch and
|
||||
# TODO(L29): 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":
|
||||
|
|
@ -284,7 +274,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
|
||||
call = self._compose_server_middleware(ctx, method, params, _inner)
|
||||
result = _dump_result(await call())
|
||||
# TODO: reject resultType values outside {"complete", "input_required"} unless the
|
||||
# TODO(L56): reject resultType values outside {"complete", "input_required"} unless the
|
||||
# corresponding extension is in this request's _meta clientCapabilities.extensions; the
|
||||
# explicit MUST-reject is client-side (basic/index.mdx ResultType), this enforces it proactively.
|
||||
if is_spec_method:
|
||||
|
|
@ -312,7 +302,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
params: Mapping[str, Any] | None,
|
||||
) -> None:
|
||||
meta = _extract_meta(params)
|
||||
version = _resolve_protocol_version(self.connection.protocol_version, meta, dctx.message_metadata)
|
||||
version = self.connection.protocol_version
|
||||
ctx = self._make_context(dctx, meta, version)
|
||||
|
||||
async def _inner() -> None:
|
||||
|
|
@ -373,7 +363,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
def _make_context(
|
||||
self, dctx: DispatchContext[TransportContext], meta: RequestParamsMeta | None, protocol_version: str
|
||||
) -> ServerRequestContext[LifespanT, Any]:
|
||||
# TODO(maxisbey): remove for Context rework. Reads the SHTTP per-request
|
||||
# TODO(L54): 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
|
||||
|
|
@ -383,8 +373,12 @@ class ServerRunner(Generic[LifespanT]):
|
|||
close_standalone_sse_stream = md.close_standalone_sse_stream
|
||||
else:
|
||||
request = close_sse_stream = close_standalone_sse_stream = None
|
||||
# Per-request session: `dctx` is the request-scoped channel (auto-threads
|
||||
# its own request_id on streamable HTTP); the standalone channel is read
|
||||
# off `connection.outbound`. `related_request_id` on the public API selects.
|
||||
session = ServerSession(dctx, self.connection)
|
||||
return ServerRequestContext(
|
||||
session=self.session,
|
||||
session=session,
|
||||
lifespan_context=self.lifespan_state,
|
||||
request_id=dctx.request_id,
|
||||
meta=meta,
|
||||
|
|
@ -405,8 +399,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
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
|
||||
opts = self.init_options if self.init_options is not None else self.server.create_initialization_options()
|
||||
return InitializeResult(
|
||||
protocol_version=negotiated,
|
||||
capabilities=opts.capabilities,
|
||||
|
|
@ -420,3 +413,82 @@ class ServerRunner(Generic[LifespanT]):
|
|||
),
|
||||
instructions=opts.instructions,
|
||||
)
|
||||
|
||||
|
||||
async def serve_connection(
|
||||
server: Server[LifespanT],
|
||||
dispatcher: Dispatcher[Any],
|
||||
*,
|
||||
connection: Connection,
|
||||
lifespan_state: LifespanT,
|
||||
init_options: InitializationOptions | None = None,
|
||||
task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
|
||||
) -> None:
|
||||
"""Drive ``dispatcher`` until the underlying channel closes.
|
||||
|
||||
The loop-mode driver: builds the kernel, hands `on_request`/`on_notify`
|
||||
to `dispatcher.run()`, and tears down `connection.exit_stack` (shielded)
|
||||
on the way out. The entry constructs the `Connection`; this only consumes
|
||||
it.
|
||||
"""
|
||||
runner = ServerRunner(server, connection, lifespan_state, init_options=init_options)
|
||||
try:
|
||||
await dispatcher.run(runner.on_request, runner.on_notify, task_status=task_status)
|
||||
finally:
|
||||
await aclose_shielded(connection)
|
||||
|
||||
|
||||
async def serve_loop(
|
||||
server: Server[LifespanT],
|
||||
read_stream: ReadStream[SessionMessage | Exception],
|
||||
write_stream: WriteStream[SessionMessage],
|
||||
*,
|
||||
lifespan_state: LifespanT,
|
||||
session_id: str | None = None,
|
||||
init_options: InitializationOptions | None = None,
|
||||
raise_exceptions: bool = False,
|
||||
) -> None:
|
||||
"""Drive ``server`` in loop mode over a stream pair until the channel closes.
|
||||
|
||||
Builds the loop-mode `JSONRPCDispatcher` + `Connection` and hands them to
|
||||
`serve_connection`, so loop-mode callers share one dispatcher-construction
|
||||
recipe (notably the `inline_methods={"initialize"}` rule). Callers that own
|
||||
a lifespan (the streamable-HTTP manager) pass it in; callers that don't
|
||||
(`Server.run` for stdio/memory) enter the lifespan and then call this.
|
||||
"""
|
||||
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: SHOULD NOT, not MUST NOT) sees the initialized
|
||||
# state instead of failing the init-gate.
|
||||
inline_methods=frozenset({"initialize"}),
|
||||
)
|
||||
connection = Connection.for_loop(dispatcher, session_id=session_id)
|
||||
await serve_connection(
|
||||
server, dispatcher, connection=connection, lifespan_state=lifespan_state, init_options=init_options
|
||||
)
|
||||
|
||||
|
||||
async def serve_one(
|
||||
server: Server[LifespanT],
|
||||
request: JSONRPCRequest,
|
||||
*,
|
||||
connection: Connection,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
lifespan_state: LifespanT,
|
||||
) -> JSONRPCResponse | JSONRPCError:
|
||||
"""Handle a single ``request`` and return its JSON-RPC reply.
|
||||
|
||||
The single-exchange driver: builds the kernel, runs `on_request` once for
|
||||
`request` under `dctx`, maps the outcome to a `JSONRPCResponse` /
|
||||
`JSONRPCError` via `to_jsonrpc_response`, and tears down
|
||||
`connection.exit_stack` (shielded) on the way out. The entry constructs
|
||||
the (born-ready) `Connection` and the `dctx`; this only consumes them.
|
||||
"""
|
||||
runner = ServerRunner(server, connection, lifespan_state)
|
||||
try:
|
||||
return await to_jsonrpc_response(request.id, runner.on_request(dctx, request.method, request.params))
|
||||
finally:
|
||||
await aclose_shielded(connection)
|
||||
|
|
|
|||
Loading…
Reference in a new issue