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

Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>

tests/server/mcpserver/test_server.py

@@ -1528,6 +1528,7 @@ async def test_report_progress_passes_related_request_id():
     request_context = ServerRequestContext(
         request_id="req-abc-123",
         session=mock_session,
+        method="tools/call",
         meta={"progress_token": "tok-1"},
         lifespan_context=None,
         protocol_version="2025-11-25",

tests/server/test_otel.py

@@ -0,0 +1,205 @@
+"""Tests for `OpenTelemetryMiddleware` (the context-tier OTel span middleware)."""
+
+from dataclasses import replace
+from typing import Any
+
+import anyio
+import pytest
+from opentelemetry.trace import SpanKind, StatusCode
+
+from mcp.server._otel import OpenTelemetryMiddleware
+from mcp.server.context import CallNext
+from mcp.server.lowlevel.server import Server
+from mcp.server.runner import otel_middleware
+from mcp.shared._otel import inject_trace_context
+from mcp.shared.exceptions import MCPError
+from mcp.types import CallToolRequestParams, ListToolsResult, NotificationParams, PaginatedRequestParams, Tool
+
+from .conftest import SpanCapture
+from .test_runner import Ctx, SrvT, connected_runner
+
+
+@pytest.fixture
+def server() -> SrvT:
+    async def list_tools(ctx: Ctx, params: PaginatedRequestParams | None) -> ListToolsResult:
+        return ListToolsResult(tools=[Tool(name="t", input_schema={"type": "object"})])
+
+    return Server(name="test-server", version="0.0.1", on_list_tools=list_tools)
+
+
+async def _ok_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
+    return {"content": [], "isError": False}
+
+
+@pytest.mark.anyio
+async def test_emits_server_span_with_method_and_target(server: SrvT, spans: SpanCapture):
+    server.add_request_handler("tools/call", CallToolRequestParams, _ok_tool)
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        result = await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
+    assert result == {"content": [], "isError": False}
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.name == "MCP handle tools/call mytool"
+    assert span.attributes is not None
+    assert span.attributes["mcp.method.name"] == "tools/call"
+    assert isinstance(span.attributes["jsonrpc.request.id"], str)
+    assert span.status.status_code == StatusCode.UNSET
+
+
+@pytest.mark.anyio
+async def test_notification_span_omits_request_id(server: SrvT, spans: SpanCapture):
+    async def on_roots(ctx: Ctx, params: NotificationParams | None) -> None:
+        return None
+
+    server.add_notification_handler("notifications/roots/list_changed", NotificationParams, on_roots)
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.notify("notifications/roots/list_changed", None)
+        await anyio.wait_all_tasks_blocked()
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.name == "MCP handle notifications/roots/list_changed"
+    assert span.attributes is not None
+    assert span.attributes["mcp.method.name"] == "notifications/roots/list_changed"
+    assert "jsonrpc.request.id" not in span.attributes
+
+
+@pytest.mark.anyio
+async def test_nests_under_ambient_span_when_no_traceparent(server: SrvT, spans: SpanCapture):
+    """With no `_meta` on the inbound message (a non-SDK client), the
+    context-tier span must parent to the ambient current span (here, the
+    dispatch-tier `otel_middleware` 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
+        # both server tiers see 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
+
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server, dispatch_middleware=[strip_meta, otel_middleware]) 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]
+    assert len(server_spans) == 2
+    [outer] = [s for s in server_spans if s.parent is None]
+    [inner] = [s for s in server_spans if s.parent is not None]
+    assert inner.context is not None and outer.context is not None
+    assert inner.parent is not None
+    assert inner.context.trace_id == outer.context.trace_id
+    assert inner.parent.span_id == outer.context.span_id
+
+
+@pytest.mark.anyio
+async def test_nests_under_ambient_span_when_meta_lacks_traceparent(server: SrvT, spans: SpanCapture):
+    """`_meta` is present but carries no `traceparent` (e.g. only a
+    `progressToken`). `extract()` would yield an empty Context here, which
+    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
+
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server, dispatch_middleware=[replace_meta, otel_middleware]) 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]
+    assert len(server_spans) == 2
+    [outer] = [s for s in server_spans if s.parent is None]
+    [inner] = [s for s in server_spans if s.parent is not None]
+    assert inner.context is not None and outer.context is not None
+    assert inner.parent is not None
+    assert inner.context.trace_id == outer.context.trace_id
+    assert inner.parent.span_id == outer.context.span_id
+
+
+@pytest.mark.anyio
+async def test_extracts_trace_context_from_meta(server: SrvT, spans: SpanCapture):
+    meta: dict[str, Any] = {}
+    inject_trace_context(meta)
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.send_raw_request("tools/list", {"_meta": meta})
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.parent is not None
+
+
+@pytest.mark.anyio
+async def test_records_error_status_on_mcp_error(server: SrvT, spans: SpanCapture):
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("resources/list", None)
+        assert exc.value.error.code != 0
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.status.status_code == StatusCode.ERROR
+    assert span.status.description == "Method not found"
+    assert not [e for e in span.events if e.name == "exception"]
+
+
+@pytest.mark.anyio
+async def test_validation_failure_sets_sanitized_status(server: SrvT, spans: SpanCapture):
+    server.add_request_handler("tools/call", CallToolRequestParams, _ok_tool)
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        with pytest.raises(MCPError):
+            await client.send_raw_request("tools/call", {"name": 123})
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.status.status_code == StatusCode.ERROR
+    assert span.status.description == "Invalid request parameters"
+    assert not span.events
+
+
+@pytest.mark.anyio
+async def test_records_error_status_on_handler_exception(server: SrvT, spans: SpanCapture):
+    async def failing(ctx: Ctx, params: PaginatedRequestParams | None) -> Any:
+        raise ValueError("handler blew up")
+
+    server.add_request_handler("tools/list", PaginatedRequestParams, failing)
+    server.middleware.append(OpenTelemetryMiddleware())
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        with pytest.raises(MCPError):
+            await client.send_raw_request("tools/list", None)
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.status.status_code == StatusCode.ERROR
+    assert span.status.description == "handler blew up"
+    [event] = [e for e in span.events if e.name == "exception"]
+    assert event.attributes is not None
+    assert event.attributes["exception.type"] == "ValueError"
+
+
+@pytest.mark.anyio
+async def test_passes_rewritten_context_through(server: SrvT, spans: SpanCapture):
+    seen_arguments: dict[str, Any] = {}
+
+    async def call_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
+        seen_arguments.update(params.arguments or {})
+        return {"content": [], "isError": False}
+
+    async def inject_arg(ctx: Ctx, call_next: CallNext) -> Any:
+        assert ctx.params is not None
+        arguments = {**ctx.params.get("arguments", {}), "injected": True}
+        return await call_next(replace(ctx, params={**ctx.params, "arguments": arguments}))
+
+    server.add_request_handler("tools/call", CallToolRequestParams, call_tool)
+    server.middleware.extend([OpenTelemetryMiddleware(), inject_arg])
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {"x": 1}})
+    assert seen_arguments == {"x": 1, "injected": True}
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.name == "MCP handle tools/call mytool"

tests/server/test_runner.py

@@ -602,9 +602,9 @@ async def test_runner_dispatch_middleware_wraps_everything_including_initialize(
 async def test_runner_server_middleware_wraps_every_request_including_initialize(server: SrvT):
     seen: list[tuple[str, Any]] = []
 
-    async def ctx_mw(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
-        seen.append((method, params))
-        return await call_next()
+    async def ctx_mw(ctx: Ctx, call_next: Any) -> Any:
+        seen.append((ctx.method, ctx.params))
+        return await call_next(ctx)
 
     server.middleware.append(ctx_mw)
     async with connected_runner(server) as (client, _):
@@ -621,9 +621,9 @@ async def test_runner_middleware_raise_after_call_next_on_initialize_leaves_conn
     client as an error and skips the state commit: the pre-init gate stays
     closed and `connection.initialized` never fires."""
 
-    async def reject_initialize(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
-        result = await call_next()
-        if method == "initialize":
+    async def reject_initialize(ctx: Ctx, call_next: Any) -> Any:
+        result = await call_next(ctx)
+        if ctx.method == "initialize":
             raise MCPError(code=INTERNAL_ERROR, message="rejected by middleware")
         return result
 
@@ -646,11 +646,11 @@ async def test_runner_middleware_raise_after_call_next_on_initialize_leaves_conn
 async def test_runner_server_middleware_observes_method_not_found_via_call_next_raise(server: SrvT):
     seen: list[tuple[str, type[BaseException] | None]] = []
 
-    async def observe(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+    async def observe(ctx: Ctx, call_next: Any) -> Any:
         try:
-            return await call_next()
+            return await call_next(ctx)
         except MCPError as e:
-            seen.append((method, type(e)))
+            seen.append((ctx.method, type(e)))
             raise
 
     server.middleware.append(observe)
@@ -668,9 +668,9 @@ async def test_runner_server_middleware_wraps_notifications(server: SrvT):
     `ctx.request_id is None`."""
     seen: list[tuple[str, bool]] = []
 
-    async def observe(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
-        seen.append((method, ctx.request_id is None))
-        return await call_next()
+    async def observe(ctx: Ctx, call_next: Any) -> Any:
+        seen.append((ctx.method, ctx.request_id is None))
+        return await call_next(ctx)
 
     async def on_roots(ctx: Ctx, params: NotificationParams | None) -> None:
         return None
@@ -713,9 +713,9 @@ async def test_runner_server_middleware_runs_outermost_first(server: SrvT):
     order: list[str] = []
 
     def make_mw(tag: str) -> Any:
-        async def mw(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        async def mw(ctx: Ctx, call_next: Any) -> Any:
             order.append(f"{tag}-in")
-            result = await call_next()
+            result = await call_next(ctx)
             order.append(f"{tag}-out")
             return result
 
@@ -761,9 +761,9 @@ async def test_runner_server_middleware_observes_handler_error_data_as_mcp_error
     successful-looking `ErrorData` return."""
     seen: list[MCPError] = []
 
-    async def observe(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+    async def observe(ctx: Ctx, call_next: Any) -> Any:
         try:
-            return await call_next()
+            return await call_next(ctx)
         except MCPError as e:
             seen.append(e)
             raise
@@ -785,7 +785,7 @@ async def test_runner_middleware_returning_error_data_produces_jsonrpc_error(ser
     """A middleware that short-circuits with an `ErrorData` return gets the
     same treatment as a handler return: the wire sees a JSON-RPC error."""
 
-    async def short_circuit(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+    async def short_circuit(ctx: Ctx, call_next: Any) -> Any:
         return ErrorData(code=INVALID_PARAMS, message="denied")
 
     server.middleware.append(short_circuit)
@@ -956,10 +956,10 @@ async def test_runner_middleware_short_circuit_on_a_wrong_version_spec_method_sk
     spec method absent at the negotiated version owns the result shape; the
     outbound sieve has no `(method, version)` row and must not raise."""
 
-    async def short_circuit(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
-        if method == "server/discover":
+    async def short_circuit(ctx: Ctx, call_next: Any) -> Any:
+        if ctx.method == "server/discover":
             return {"ok": True}
-        return await call_next()
+        return await call_next(ctx)
 
     server.middleware.append(short_circuit)
     async with connected_runner(server) as (client, runner):
@@ -1047,13 +1047,12 @@ async def test_otel_trace_context_propagates_client_to_server(server: SrvT, span
 
 @pytest.mark.anyio
 async def test_otel_middleware_malformed_traceparent_degrades_to_no_parent(server: SrvT, spans: SpanCapture):
-    """A non-string traceparent in `_meta` must not fail the request; the
-    server span simply gets no parent."""
+    """A non-string traceparent in `_meta` must not fail the request; the server span simply gets no parent."""
 
-    def break_traceparent(next_on_request: OnRequest) -> OnRequest:
-        async def wrapped(dctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
+    def break_traceparent(call_next: OnRequest) -> OnRequest:
+        async def wrapped(ctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
             mangled = {"_meta": {"traceparent": 123}} if method == "tools/list" else params
-            return await next_on_request(dctx, method, mangled)
+            return await call_next(ctx, method, mangled)
 
         return wrapped