"""`ServerRunner` - the per-connection handler kernel. `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. """ from __future__ import annotations import logging 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 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._stream_protocols import ReadStream, WriteStream from mcp.shared.dispatcher import DispatchContext, Dispatcher, DispatchMiddleware, OnNotify, OnRequest from mcp.shared.exceptions import MCPError 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 ( INTERNAL_ERROR, INVALID_PARAMS, LATEST_PROTOCOL_VERSION, METHOD_NOT_FOUND, ErrorData, Implementation, InitializeRequestParams, InitializeResult, JSONRPCError, JSONRPCRequest, JSONRPCResponse, RequestId, RequestParams, RequestParamsMeta, ) from mcp.types import methods as _methods if TYPE_CHECKING: from mcp.server.lowlevel.server import Server __all__ = [ "CallNext", "ServerMiddleware", "ServerRunner", "aclose_shielded", "otel_middleware", "serve_connection", "serve_loop", "serve_one", "to_jsonrpc_response", ] logger = logging.getLogger(__name__) LifespanT = TypeVar("LifespanT", default=Any) _INIT_EXEMPT: frozenset[str] = frozenset({"ping"}) _EXIT_STACK_CLOSE_TIMEOUT: float = 5 """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: """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 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 []"`, `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") 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 handler kernel. One instance per client connection.""" server: Server[LifespanT] connection: Connection lifespan_state: LifespanT _: KW_ONLY init_options: InitializationOptions | None = None """`InitializeResult` payload. Defaults to `server.create_initialization_options()`.""" dispatch_middleware: Sequence[DispatchMiddleware] = (otel_middleware,) @cached_property def on_request(self) -> OnRequest: """`_on_request` wrapped in `dispatch_middleware`, outermost-first. 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], method: str, params: Mapping[str, Any] | None, ) -> dict[str, Any]: meta = _extract_meta(params) version = self.connection.protocol_version ctx = self._make_context(dctx, meta, version) is_spec_method = method in _methods.SPEC_CLIENT_METHODS async def _inner() -> HandlerResult: # Pinned compat: spec methods are surface-validated before lookup, # so malformed params are INVALID_PARAMS even with no handler # registered. Custom methods miss the monolith map and fall through # to `entry.params_type` exactly as before. if is_spec_method: try: _methods.validate_client_request(method, version, params) except KeyError: raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method) from None # 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": return self._handle_initialize(params) # Methods without a handler are METHOD_NOT_FOUND regardless of # initialization state: JSON-RPC 2.0 reserves -32601 for "not # available on this server", and clients probing a server before # the handshake key off that code. The init gate below therefore # only ever applies to methods the server actually serves. entry = self.server.get_request_handler(method) if entry is None: raise MCPError(code=METHOD_NOT_FOUND, message="Method not found", data=method) 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="") # 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()) # 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: try: result = _methods.serialize_server_result(method, version, result) except KeyError: # Middleware short-circuited a wrong-version spec method without # calling `call_next`; it owns the result shape. pass except ValidationError: # Server bug, not client fault. Detail stays in the server log: # pydantic messages echo the result body. logger.exception("handler for %r returned an invalid result", method) raise MCPError(code=INTERNAL_ERROR, message="Handler returned an invalid result") from None 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: meta = _extract_meta(params) version = self.connection.protocol_version ctx = self._make_context(dctx, meta, version) async def _inner() -> None: if method in _methods.SPEC_CLIENT_NOTIFICATION_METHODS: try: _methods.validate_client_notification(method, version, params) except KeyError: logger.debug("dropped %r: not defined at %s", method, version) return except ValidationError: logger.warning("dropped %r: malformed params", method) return if method == "notifications/initialized": # Surface validation above already rejected a malformed body, so # commit; fall through so a registered handler observes an # initialized connection. 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, protocol_version: str ) -> ServerRequestContext[LifespanT, Any]: # 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 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 # 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=session, lifespan_context=self.lifespan_state, request_id=dctx.request_id, meta=meta, protocol_version=protocol_version, 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) 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, 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, ) 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)