modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Slim ServerMiddleware to (ctx, call_next) and add OpenTelemetryMiddleware (#2941)

Browse source 
Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>
This commit is contained in:
Marcelo Trylesinski 2026-06-22 15:46:30 +02:00 • committed by GitHub
commit ad81ca234a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 374 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

51
src/mcp/server/_otel.py Normal file
View file

@ -0,0 +1,51 @@
from __future__ import annotations
from typing import Any
from opentelemetry.trace import SpanKind, StatusCode
from pydantic import ValidationError
from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
from mcp.shared._otel import extract_trace_context, otel_span
from mcp.shared.exceptions import MCPError
class OpenTelemetryMiddleware(ServerMiddleware[Any]):
"""Context-tier middleware that wraps each inbound message in an OpenTelemetry span.
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. Requests and notifications both get a span;
`jsonrpc.request.id` is set only when `ctx.request_id` is present (notifications
have none).
"""
async def __call__(self, ctx: ServerRequestContext[Any, Any], call_next: CallNext) -> HandlerResult:
name = ctx.params.get("name") if ctx.params else None
target = name if isinstance(name, str) else None
attributes: dict[str, Any] = {"mcp.method.name": ctx.method}
if ctx.request_id is not None:
attributes["jsonrpc.request.id"] = str(ctx.request_id)
with otel_span(
name=f"MCP handle {ctx.method}{f' {target}' if target else ''}",
kind=SpanKind.SERVER,
attributes=attributes,
context=extract_trace_context(ctx.meta),
record_exception=False,
set_status_on_exception=False,
) as span:
try:
return await call_next(ctx)
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

37
src/mcp/server/context.py
View file

@ -15,7 +15,7 @@ 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.
# Invariant: parametrizes a mutable dataclass field; dict default matches the default lifespan.
LifespanContextT = TypeVar("LifespanContextT", default=dict[str, Any])
RequestT = TypeVar("RequestT", default=Any)
@ -33,6 +33,8 @@ class ServerRequestContext(Generic[LifespanContextT, RequestT]):
session: ServerSession
lifespan_context: LifespanContextT
protocol_version: str
method: str
params: Mapping[str, Any] | None = None
request_id: RequestId | None = None
meta: RequestParamsMeta | None = None
request: RequestT | None = None
@ -113,39 +115,44 @@ 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]]
CallNext = Callable[["ServerRequestContext[Any, Any]"], Awaitable[HandlerResult]]
"""Invokes the rest of the chain. Pass the `ctx` through; rewrite `method` or
`params` with `dataclasses.replace(ctx, ...)` to alter what the handler sees."""
_MwLifespanT = TypeVar("_MwLifespanT")
class ServerMiddleware(Protocol[_MwLifespanT]):
"""Context-tier middleware: `(ctx, method, params, call_next) -> result`.
"""Context-tier middleware: `(ctx, 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/initialized` all run inside `call_next(ctx)`.
`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`.
The method and the raw inbound params are `ctx.method` and `ctx.params` (no
model validation has happened yet). To rewrite either before the handler
runs, pass an adjusted context: `await call_next(replace(ctx, params=...))`.
`ctx.request_id is None` distinguishes a notification from a request. For
notifications `call_next()` returns `None` (a dropped or unhandled
notifications `call_next(ctx)` 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.
!!! 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.
`initialize` is observed but not rewritable: the post-chain handshake
commit reads the wire params, so to veto the handshake raise *before*
`call_next()`.
`Server[L].middleware` holds `ServerMiddleware[L]`, so an app-specific
middleware sees `ctx.lifespan_context: L`. While the context is the
@ -162,7 +169,5 @@ class ServerMiddleware(Protocol[_MwLifespanT]):
async def __call__(
self,
ctx: ServerRequestContext[_MwLifespanT, Any],
method: str,
params: Mapping[str, Any] | None,
call_next: CallNext,
) -> HandlerResult: ...

2
src/mcp/server/lowlevel/server.py
View file

@ -225,7 +225,7 @@ class Server(Generic[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`.
# `(ctx, call_next)`. Applied in `ServerRunner._on_request`.
# TODO(L54): provisional - signature and semantics change with the
# Context/middleware rework (covariant `Context[L]`, outbound seam) before
# v2 final.

67
src/mcp/server/runner.py
View file

@ -93,7 +93,7 @@ def _extract_meta(params: Mapping[str, Any] | None) -> RequestParamsMeta | None:
return None
def otel_middleware(next_on_request: OnRequest) -> OnRequest:
def otel_middleware(call_next: 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
@ -129,7 +129,7 @@ def otel_middleware(next_on_request: OnRequest) -> OnRequest:
set_status_on_exception=False,
) as span:
try:
return await next_on_request(dctx, method, params)
return await call_next(dctx, method, params)
except MCPError as e:
span.set_status(StatusCode.ERROR, e.error.message)
raise
@ -200,6 +200,14 @@ async def to_jsonrpc_response(request_id: RequestId, coro: Awaitable[dict[str, A
return JSONRPCResponse(jsonrpc="2.0", id=request_id, result=result)
def _apply_middleware(
middleware: ServerMiddleware[Any], call_next: CallNext, ctx: ServerRequestContext[Any, Any]
) -> Awaitable[HandlerResult]:
"""Adapt one middleware to the `CallNext` shape: bind `call_next`, take
`ctx` at call time so a rewritten context flows down the chain."""
return middleware(ctx, call_next)
@dataclass
class ServerRunner(Generic[LifespanT]):
"""Per-connection handler kernel. One instance per client connection."""
@ -220,7 +228,9 @@ class ServerRunner(Generic[LifespanT]):
wraps everything - initialize, METHOD_NOT_FOUND, validation failures
included.
"""
return reduce(lambda h, mw: mw(h), reversed(self.dispatch_middleware), self._on_request)
return reduce(
lambda handler, middleware: middleware(handler), reversed(self.dispatch_middleware), self._on_request
)
@cached_property
def on_notify(self) -> OnNotify:
@ -234,15 +244,18 @@ class ServerRunner(Generic[LifespanT]):
) -> dict[str, Any]:
meta = _extract_meta(params)
version = self.connection.protocol_version
ctx = self._make_context(dctx, meta, version)
ctx = self._make_context(dctx, method, params, meta, version)
is_spec_method = method in _methods.SPEC_CLIENT_METHODS
async def _inner() -> HandlerResult:
async def _inner(ctx: ServerRequestContext[LifespanT, Any]) -> HandlerResult:
# Read method/params off `ctx` so a middleware that rewrote them via
# `call_next(replace(ctx, ...))` reaches lookup and the handler.
method, params = ctx.method, ctx.params
# 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:
if method in _methods.SPEC_CLIENT_METHODS:
try:
_methods.validate_client_request(method, version, params)
except KeyError:
@ -272,8 +285,8 @@ class ServerRunner(Generic[LifespanT]):
raise MCPError.from_error_data(result)
return result
call = self._compose_server_middleware(ctx, method, params, _inner)
result = _dump_result(await call())
call = self._compose_server_middleware(_inner)
result = _dump_result(await call(ctx))
# 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.
@ -292,6 +305,11 @@ class ServerRunner(Generic[LifespanT]):
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.
# TODO: this re-reads the wire `params`, so a middleware that rewrote
# `ctx.params` (or `ctx.method`, or short-circuited without `call_next`)
# can leave `connection.protocol_version` out of step with the
# `InitializeResult` `_inner` produced. Resolve when `initialize` becomes
# a built-in handler so commit and result derive from one negotiation.
self.connection.client_params, self.connection.protocol_version = self._negotiate_initialize(params)
return result
@ -303,9 +321,10 @@ class ServerRunner(Generic[LifespanT]):
) -> None:
meta = _extract_meta(params)
version = self.connection.protocol_version
ctx = self._make_context(dctx, meta, version)
ctx = self._make_context(dctx, method, params, meta, version)
async def _inner() -> None:
async def _inner(ctx: ServerRequestContext[LifespanT, Any]) -> None:
method, params = ctx.method, ctx.params
if method in _methods.SPEC_CLIENT_NOTIFICATION_METHODS:
try:
_methods.validate_client_notification(method, version, params)
@ -335,33 +354,33 @@ class ServerRunner(Generic[LifespanT]):
return
await entry.handler(ctx, typed_params)
call = self._compose_server_middleware(ctx, method, params, _inner)
call = self._compose_server_middleware(_inner)
try:
await call()
await call(ctx)
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:
def _compose_server_middleware(self, 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.
observes every inbound message. The composed callable takes the `ctx`
at call time, so a middleware can rewrite it for the rest of the chain.
"""
call = inner
for mw in reversed(self.server.middleware):
call = partial(mw, ctx, method, params, call)
for middleware in reversed(self.server.middleware):
call = partial(_apply_middleware, middleware, call)
return call
def _make_context(
self, dctx: DispatchContext[TransportContext], meta: RequestParamsMeta | None, protocol_version: str
self,
dctx: DispatchContext[TransportContext],
method: str,
params: Mapping[str, Any] | None,
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
@ -380,6 +399,8 @@ class ServerRunner(Generic[LifespanT]):
return ServerRequestContext(
session=session,
lifespan_context=self.lifespan_state,
method=method,
params=params,
request_id=dctx.request_id,
meta=meta,
protocol_version=protocol_version,