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
5 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,

26
src/mcp/shared/_otel.py
View file

@ -2,13 +2,14 @@
from __future__ import annotations
from collections.abc import Iterator
from collections.abc import Generator, Mapping
from contextlib import contextmanager
from typing import Any
from opentelemetry.context import Context
from opentelemetry.propagate import extract, inject
from opentelemetry.trace import SpanKind, get_tracer
from opentelemetry.trace import SpanKind, get_current_span, get_tracer
from opentelemetry.trace.span import Span
_tracer = get_tracer("mcp-python-sdk")
@ -22,7 +23,7 @@ def otel_span(
context: Context | None = None,
record_exception: bool = True,
set_status_on_exception: bool = True,
) -> Iterator[Any]:
) -> Generator[Span]:
"""Create an OTel span."""
with _tracer.start_as_current_span(
name,
@ -40,13 +41,20 @@ def inject_trace_context(meta: dict[str, Any]) -> None:
inject(meta)
def extract_trace_context(meta: dict[str, Any]) -> Context | None:
def extract_trace_context(meta: Mapping[str, Any] | None) -> 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.
Returns `None` when the carrier is absent, malformed, or carries no
valid `traceparent`, so callers fall through to ambient parenting; an
explicit empty `Context` would orphan the span instead of nesting under
the current one.
"""
try:
return extract(meta)
except (TypeError, ValueError):
if not meta:
return None
try:
ctx = extract(meta)
except (ValueError, TypeError):
return None
if not get_current_span(ctx).get_span_context().is_valid:
return None
return ctx