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

Remove the dispatch-tier middleware hook (#2997)

Browse source
This commit is contained in:
Marcelo Trylesinski 2026-06-26 17:08:16 +02:00 • committed by GitHub
commit 3945bdde11
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 32 additions and 79 deletions
Download patch file Download diff file Expand all files Collapse all files

41
tests/server/test_otel.py
View file

@ -5,6 +5,7 @@ so these tests assert against the default-configured server rather than appendin
the middleware by hand.
"""
from collections.abc import Callable
from dataclasses import replace
from typing import Any
@ -20,6 +21,7 @@ from mcp_types import (
ListToolsResult,
NotificationParams,
PaginatedRequestParams,
RequestParamsMeta,
Tool,
)
from opentelemetry.trace import SpanKind, StatusCode
@ -166,16 +168,17 @@ async def test_notification_span_omits_request_id(server: SrvT, spans: SpanCaptu
assert "jsonrpc.request.id" not in span.attributes
def _ambient_span(call_next: Any) -> Any:
"""Dispatch-tier wrapper that opens an ambient SERVER span around the whole
request, so the context-tier span has a current span to nest under when the
inbound message carries no traceparent."""
def _ambient(rewrite_meta: Callable[[RequestParamsMeta | None], RequestParamsMeta | None]) -> Any:
"""A middleware placed outside `OpenTelemetryMiddleware` (head of
`Server.middleware`) that opens an ambient SERVER span around the request
and rewrites `ctx.meta` via `rewrite_meta`, so the context-tier span sees
the no-traceparent path yet has a current span to nest under."""
async def wrapped(dctx: Any, method: str, params: dict[str, Any] | None) -> Any:
async def middleware(ctx: Ctx, call_next: CallNext) -> Any:
with otel_span("ambient", kind=SpanKind.SERVER):
return await call_next(dctx, method, params)
return await call_next(replace(ctx, meta=rewrite_meta(ctx.meta)))
return wrapped
return middleware
@pytest.mark.anyio
@ -184,16 +187,10 @@ async def test_nests_under_ambient_span_when_no_traceparent(server: SrvT, spans:
parent to the ambient current span rather than become an orphan root.
SDK-defined: SEP-414 only covers the traceparent-present case."""
def strip_meta(call_next: Any) -> Any:
# The in-process client always injects `_meta.traceparent`; strip it so
# the span sees the no-carrier path.
async def wrapped(dctx: Any, method: str, params: dict[str, Any] | None) -> Any:
stripped = {k: v for k, v in (params or {}).items() if k != "_meta"}
return await call_next(dctx, method, stripped or None)
return wrapped
async with connected_runner(server, dispatch_middleware=[_ambient_span, strip_meta]) as (client, _):
# The in-process client always injects `_meta.traceparent`; drop it so the
# span sees the no-carrier path.
server.middleware.insert(0, _ambient(lambda _meta: None))
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("tools/list", None)
server_spans = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
@ -213,14 +210,8 @@ async def test_nests_under_ambient_span_when_meta_lacks_traceparent(server: SrvT
would orphan the span; the middleware must fall through to ambient
parenting just as if `_meta` were absent."""
def replace_meta(call_next: Any) -> Any:
async def wrapped(dctx: Any, method: str, params: dict[str, Any] | None) -> Any:
rewritten = {**(params or {}), "_meta": {"progressToken": "tok"}}
return await call_next(dctx, method, rewritten)
return wrapped
async with connected_runner(server, dispatch_middleware=[_ambient_span, replace_meta]) as (client, _):
server.middleware.insert(0, _ambient(lambda _meta: {"progressToken": "tok"}))
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("tools/list", None)
server_spans = [s for s in spans.finished() if s.kind == SpanKind.SERVER]

46
tests/server/test_runner.py
View file

@ -9,7 +9,7 @@ behaviour under test. Driver tests (`serve_connection`, `serve_one`,
from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from functools import partial
from typing import Any, cast
@ -48,7 +48,7 @@ from mcp.server.runner import (
serve_one,
)
from mcp.server.session import ServerSession
from mcp.shared.dispatcher import CallOptions, DispatchContext, DispatchMiddleware, OnRequest
from mcp.shared.dispatcher import CallOptions
from mcp.shared.exceptions import MCPError
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
from mcp.shared.message import MessageMetadata
@ -91,7 +91,6 @@ async def connected_runner(
*,
initialized: bool = True,
init_options: InitializationOptions | None = None,
dispatch_middleware: list[DispatchMiddleware] | None = None,
connection: Connection | None = None,
) -> AsyncIterator[tuple[JSONRPCDispatcher[TransportContext], ServerRunner[dict[str, Any]]]]:
"""Yield `(client, runner)` running over an in-memory JSON-RPC dispatcher pair.
@ -116,7 +115,6 @@ async def connected_runner(
connection=connection,
lifespan_state={},
init_options=init_options,
dispatch_middleware=dispatch_middleware or [],
)
c_req, c_notify = echo_handlers(Recorder())
body_exc: BaseException | None = None
@ -511,8 +509,8 @@ async def test_runner_absent_wire_params_reaches_request_handler_as_defaults_mod
"""A request with no `params` member on the wire reaches the handler as
the params model with its defaults, never `None`.
The in-SDK client always attaches `_meta`, so a dispatch middleware
forwards `params=None` to model what an external client sends.
The in-SDK client always attaches `_meta`, so a middleware rewrites
`ctx.params` to `None` to model what an external client sends.
"""
seen: list[PaginatedRequestParams | None] = []
@ -520,14 +518,12 @@ async def test_runner_absent_wire_params_reaches_request_handler_as_defaults_mod
seen.append(params)
return ListToolsResult(tools=[])
def drop_params(next_on_request: OnRequest) -> OnRequest:
async def wrapped(dctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
return await next_on_request(dctx, method, None if method == "tools/list" else params)
return wrapped
async def drop_params(ctx: Ctx, call_next: Any) -> Any:
return await call_next(replace(ctx, params=None) if ctx.method == "tools/list" else ctx)
server: SrvT = Server(name="s", on_list_tools=list_tools)
async with connected_runner(server, dispatch_middleware=[drop_params]) as (client, _):
server.middleware.append(drop_params)
async with connected_runner(server) as (client, _):
await client.send_raw_request("tools/list", None)
assert seen == [PaginatedRequestParams()]
@ -543,15 +539,13 @@ async def test_runner_absent_wire_params_for_required_params_custom_method_is_in
async def greet(ctx: Ctx, params: GreetParams) -> dict[str, Any]:
raise NotImplementedError
def drop_params(next_on_request: OnRequest) -> OnRequest:
async def wrapped(dctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
return await next_on_request(dctx, method, None if method == "custom/greet" else params)
return wrapped
async def drop_params(ctx: Ctx, call_next: Any) -> Any:
return await call_next(replace(ctx, params=None) if ctx.method == "custom/greet" else ctx)
server: SrvT = Server(name="s")
server.add_request_handler("custom/greet", GreetParams, greet)
async with connected_runner(server, dispatch_middleware=[drop_params]) as (client, _):
server.middleware.append(drop_params)
async with connected_runner(server) as (client, _):
with pytest.raises(MCPError) as exc:
await client.send_raw_request("custom/greet", {"name": "x"})
assert exc.value.error.code == INVALID_PARAMS
@ -574,22 +568,6 @@ async def test_runner_on_notify_drops_before_init_and_unknown_methods(server: Sr
assert seen == [NotificationParams()] # only the post-init one reached the handler
@pytest.mark.anyio
async def test_runner_dispatch_middleware_wraps_everything_including_initialize(server: SrvT):
seen_methods: list[str] = []
def trace_mw(next_on_request: Any) -> Any:
async def wrapped(dctx: Any, method: str, params: Any) -> Any:
seen_methods.append(method)
return await next_on_request(dctx, method, params)
return wrapped
async with connected_runner(server, dispatch_middleware=[trace_mw]) as (client, _):
await client.send_raw_request("tools/list", None)
assert seen_methods == ["initialize", "tools/list"]
@pytest.mark.anyio
async def test_runner_server_middleware_wraps_every_request_including_initialize(server: SrvT):
seen: list[tuple[str, Any]] = []