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

Make OpenTelemetry tracing the single default middleware (#2995)

Browse source
This commit is contained in:
Marcelo Trylesinski 2026-06-26 15:47:37 +02:00 • committed by GitHub
commit b31d95a429
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 260 additions and 238 deletions
Download patch file Download diff file Expand all files Collapse all files

76
tests/server/test_otel.py
View file

@ -1,4 +1,9 @@
"""Tests for `OpenTelemetryMiddleware` (the context-tier OTel span middleware)."""
"""Tests for `OpenTelemetryMiddleware` (the context-tier OTel span middleware).
Every `Server` ships `OpenTelemetryMiddleware` at the head of `Server.middleware`,
so these tests assert against the default-configured server rather than appending
the middleware by hand.
"""
from dataclasses import replace
from typing import Any
@ -6,6 +11,7 @@ from typing import Any
import anyio
import pytest
from mcp_types import (
INTERNAL_ERROR,
INVALID_PARAMS,
CallToolRequestParams,
CallToolResult,
@ -21,8 +27,7 @@ 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._otel import inject_trace_context, otel_span
from mcp.shared.exceptions import MCPError
from .conftest import SpanCapture
@ -41,10 +46,14 @@ async def _ok_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
return {"content": [], "isError": False}
def test_server_ships_opentelemetry_middleware_by_default() -> None:
server = Server(name="test-server", version="0.0.1")
assert any(isinstance(m, OpenTelemetryMiddleware) for m in server.middleware)
@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": {}})
@ -65,7 +74,6 @@ async def test_tool_error_dict_result_sets_error_type(server: SrvT, spans: SpanC
return {"content": [], "isError": True}
server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
server.middleware.append(OpenTelemetryMiddleware())
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
@ -81,7 +89,6 @@ async def test_tool_error_model_result_sets_error_type(server: SrvT, spans: Span
return CallToolResult(content=[], is_error=True)
server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
server.middleware.append(OpenTelemetryMiddleware())
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
@ -99,7 +106,6 @@ async def test_snake_case_dict_result_is_not_a_tool_error(server: SrvT, spans: S
return {"content": [], "is_error": True}
server.add_request_handler("tools/call", CallToolRequestParams, err_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": {}})
@ -116,7 +122,6 @@ async def test_named_non_tool_prompt_method_omits_gen_ai_attrs(server: SrvT, spa
return {"content": [], "isError": False}
server.add_request_handler("custom/op", CallToolRequestParams, custom)
server.middleware.append(OpenTelemetryMiddleware())
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("custom/op", {"name": "thing", "arguments": {}})
@ -134,7 +139,6 @@ async def test_prompt_get_sets_prompt_name(server: SrvT, spans: SpanCapture):
return GetPromptResult(messages=[])
server.add_request_handler("prompts/get", GetPromptRequestParams, get_prompt)
server.middleware.append(OpenTelemetryMiddleware())
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("prompts/get", {"name": "myprompt"})
@ -151,7 +155,6 @@ async def test_notification_span_omits_request_id(server: SrvT, spans: SpanCaptu
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)
@ -163,24 +166,34 @@ 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."""
async def wrapped(dctx: Any, method: str, params: dict[str, Any] | None) -> Any:
with otel_span("ambient", kind=SpanKind.SERVER):
return await call_next(dctx, method, params)
return wrapped
@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.
"""With no `_meta` on the inbound message (a non-SDK client), the span must
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
# both server tiers see the no-carrier path.
# 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
server.middleware.append(OpenTelemetryMiddleware())
async with connected_runner(server, dispatch_middleware=[strip_meta, otel_middleware]) as (client, _):
async with connected_runner(server, dispatch_middleware=[_ambient_span, strip_meta]) 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]
@ -207,8 +220,7 @@ async def test_nests_under_ambient_span_when_meta_lacks_traceparent(server: SrvT
return wrapped
server.middleware.append(OpenTelemetryMiddleware())
async with connected_runner(server, dispatch_middleware=[replace_meta, otel_middleware]) as (client, _):
async with connected_runner(server, dispatch_middleware=[_ambient_span, replace_meta]) 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]
@ -225,7 +237,6 @@ async def test_nests_under_ambient_span_when_meta_lacks_traceparent(server: SrvT
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})
@ -235,7 +246,6 @@ async def test_extracts_trace_context_from_meta(server: SrvT, spans: SpanCapture
@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:
@ -253,7 +263,6 @@ async def test_records_error_status_on_mcp_error(server: SrvT, spans: SpanCaptur
@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):
@ -275,7 +284,6 @@ async def test_records_error_status_on_handler_exception(server: SrvT, spans: Sp
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):
@ -290,6 +298,28 @@ async def test_records_error_status_on_handler_exception(server: SrvT, spans: Sp
assert event.attributes["exception.type"] == "ValueError"
@pytest.mark.anyio
async def test_records_error_status_on_malformed_spec_result(server: SrvT, spans: SpanCapture):
"""Result serialization runs inside the span, so a handler returning a
malformed dict for a spec method (INTERNAL_ERROR on the wire) is recorded
on the span rather than closing it as a success."""
async def bad_result(ctx: Ctx, params: PaginatedRequestParams | None) -> dict[str, Any]:
return {"tools": 42}
server.add_request_handler("tools/list", PaginatedRequestParams, bad_result)
async with connected_runner(server) as (client, _):
spans.clear()
with pytest.raises(MCPError) as exc:
await client.send_raw_request("tools/list", None)
assert exc.value.error.code == INTERNAL_ERROR
[span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
assert span.status.status_code == StatusCode.ERROR
assert span.attributes is not None
assert span.attributes["error.type"] == str(INTERNAL_ERROR)
assert span.attributes["rpc.response.status_code"] == str(INTERNAL_ERROR)
@pytest.mark.anyio
async def test_passes_rewritten_context_through(server: SrvT, spans: SpanCapture):
seen_arguments: dict[str, Any] = {}
@ -304,7 +334,7 @@ async def test_passes_rewritten_context_through(server: SrvT, spans: SpanCapture
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])
server.middleware.append(inject_arg)
async with connected_runner(server) as (client, _):
spans.clear()
await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {"x": 1}})

106
tests/server/test_runner.py
View file

@ -21,7 +21,6 @@ from mcp_types import (
INVALID_PARAMS,
LATEST_PROTOCOL_VERSION,
METHOD_NOT_FOUND,
CallToolRequestParams,
ClientCapabilities,
ErrorData,
Implementation,
@ -35,7 +34,6 @@ from mcp_types import (
Tool,
)
from mcp_types.version import LATEST_HANDSHAKE_VERSION, LATEST_MODERN_VERSION, OLDEST_SUPPORTED_VERSION
from opentelemetry.trace import SpanKind, StatusCode
import mcp.server.runner
from mcp.server.connection import Connection
@ -46,7 +44,6 @@ from mcp.server.runner import (
ServerRunner,
_extract_meta,
aclose_shielded,
otel_middleware,
serve_connection,
serve_one,
)
@ -60,7 +57,6 @@ from mcp.shared.transport_context import TransportContext
from ..shared.conftest import jsonrpc_pair
from ..shared.test_dispatcher import Recorder, echo_handlers
from .conftest import SpanCapture
Ctx = ServerRequestContext[dict[str, Any], Any]
@ -1005,108 +1001,6 @@ async def test_runner_initialize_echoes_supported_version_and_falls_back_to_late
assert result["protocolVersion"] == LATEST_HANDSHAKE_VERSION
@pytest.mark.anyio
async def test_otel_middleware_emits_server_span_with_method_and_target(server: SrvT, spans: SpanCapture):
async def call_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
return {"content": [], "isError": False}
server.add_request_handler("tools/call", CallToolRequestParams, call_tool)
async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
spans.clear()
result = await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
assert result == {"content": [], "isError": False}
finished = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
[span] = finished
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_otel_trace_context_propagates_client_to_server(server: SrvT, spans: SpanCapture):
"""The client dispatcher injects traceparent into `_meta`; the server's
`otel_middleware` extracts it, so client and server spans share a trace."""
async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
spans.clear()
await client.send_raw_request("tools/list", None)
[client_span] = [s for s in spans.finished() if s.kind == SpanKind.CLIENT]
[server_span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
assert server_span.parent is not None
assert client_span.context is not None and server_span.context is not None
assert server_span.parent.span_id == client_span.context.span_id
assert server_span.context.trace_id == client_span.context.trace_id
assert client_span.attributes is not None and server_span.attributes is not None
assert client_span.attributes["jsonrpc.request.id"] == server_span.attributes["jsonrpc.request.id"]
@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."""
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 call_next(ctx, method, mangled)
return wrapped
async with connected_runner(server, dispatch_middleware=[break_traceparent, otel_middleware]) as (client, _):
spans.clear()
await client.send_raw_request("tools/list", None)
[server_span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
assert server_span.parent is None
@pytest.mark.anyio
async def test_otel_middleware_validation_failure_sets_sanitized_status(server: SrvT, spans: SpanCapture):
"""Malformed params set the sanitized wire message as span status and do
not record the pydantic exception (it carries client input)."""
async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
spans.clear()
with pytest.raises(MCPError) as exc:
await client.send_raw_request("tools/call", {"name": 123})
assert exc.value.error.code == INVALID_PARAMS
[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_otel_middleware_records_error_status_on_mcp_error(server: SrvT, spans: SpanCapture):
async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
spans.clear()
with pytest.raises(MCPError) as exc:
await client.send_raw_request("resources/list", None)
assert exc.value.error.code == METHOD_NOT_FOUND
[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"
# MCPError is a protocol-level response, not a crash - no traceback event.
assert not [e for e in span.events if e.name == "exception"]
@pytest.mark.anyio
async def test_otel_middleware_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)
async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
spans.clear()
with pytest.raises(MCPError) as exc:
await client.send_raw_request("tools/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 == "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_runner_connection_exit_stack_unwinds_after_run_returns(server: SrvT) -> None:
"""`runner.connection.exit_stack` is closed when the dispatcher loop ends."""