Server-side 2026-07-28 stateless support: classifier, driver split, server/discover (#2928)

This commit is contained in:
Max 2026-06-21 19:34:17 +01:00 • committed by GitHub
commit 2397319a68
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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)