chore: import upstream snapshot with attribution

tests/server/test_otel.py

@@ -0,0 +1,334 @@
+"""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 collections.abc import Callable
+from dataclasses import replace
+from typing import Any
+
+import anyio
+import pytest
+from mcp_types import (
+    INTERNAL_ERROR,
+    INVALID_PARAMS,
+    CallToolRequestParams,
+    CallToolResult,
+    GetPromptRequestParams,
+    GetPromptResult,
+    ListToolsResult,
+    NotificationParams,
+    PaginatedRequestParams,
+    RequestParamsMeta,
+    Tool,
+)
+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.shared._otel import inject_trace_context, otel_span
+from mcp.shared.exceptions import MCPError
+
+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}
+
+
+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)
+    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 == "tools/call mytool"
+    assert span.attributes is not None
+    assert span.attributes["mcp.method.name"] == "tools/call"
+    assert span.attributes["gen_ai.operation.name"] == "execute_tool"
+    assert span.attributes["gen_ai.tool.name"] == "mytool"
+    assert isinstance(span.attributes["jsonrpc.request.id"], str)
+    assert span.status.status_code == StatusCode.UNSET
+
+
+@pytest.mark.anyio
+async def test_tool_error_dict_result_sets_error_type(server: SrvT, spans: SpanCapture):
+    async def err_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
+        return {"content": [], "isError": True}
+
+    server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.attributes is not None
+    assert span.attributes["error.type"] == "tool_error"
+    assert span.status.status_code == StatusCode.ERROR
+
+
+@pytest.mark.anyio
+async def test_tool_error_model_result_sets_error_type(server: SrvT, spans: SpanCapture):
+    async def err_tool(ctx: Ctx, params: CallToolRequestParams) -> CallToolResult:
+        return CallToolResult(content=[], is_error=True)
+
+    server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.attributes is not None
+    assert span.attributes["error.type"] == "tool_error"
+    assert span.status.status_code == StatusCode.ERROR
+
+
+@pytest.mark.anyio
+async def test_snake_case_dict_result_is_not_a_tool_error(server: SrvT, spans: SpanCapture):
+    # `is_error` is alias-only on the wire, so serialization drops it; the result reaches the
+    # client as a success and the span must not contradict that.
+    async def err_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
+        return {"content": [], "is_error": True}
+
+    server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        result = await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
+    assert result == {"content": []}
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.attributes is not None
+    assert "error.type" not in span.attributes
+    assert span.status.status_code == StatusCode.UNSET
+
+
+@pytest.mark.anyio
+async def test_named_non_tool_prompt_method_omits_gen_ai_attrs(server: SrvT, spans: SpanCapture):
+    async def custom(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
+        return {"content": [], "isError": False}
+
+    server.add_request_handler("custom/op", CallToolRequestParams, custom)
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.send_raw_request("custom/op", {"name": "thing", "arguments": {}})
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.name == "custom/op thing"
+    assert span.attributes is not None
+    assert "gen_ai.operation.name" not in span.attributes
+    assert "gen_ai.tool.name" not in span.attributes
+    assert "gen_ai.prompt.name" not in span.attributes
+
+
+@pytest.mark.anyio
+async def test_prompt_get_sets_prompt_name(server: SrvT, spans: SpanCapture):
+    async def get_prompt(ctx: Ctx, params: GetPromptRequestParams) -> GetPromptResult:
+        return GetPromptResult(messages=[])
+
+    server.add_request_handler("prompts/get", GetPromptRequestParams, get_prompt)
+    async with connected_runner(server) as (client, _):
+        spans.clear()
+        await client.send_raw_request("prompts/get", {"name": "myprompt"})
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.name == "prompts/get myprompt"
+    assert span.attributes is not None
+    assert span.attributes["gen_ai.prompt.name"] == "myprompt"
+    assert "gen_ai.operation.name" not in span.attributes
+
+
+@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)
+    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 == "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
+
+
+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 middleware(ctx: Ctx, call_next: CallNext) -> Any:
+        with otel_span("ambient", kind=SpanKind.SERVER):
+            return await call_next(replace(ctx, meta=rewrite_meta(ctx.meta)))
+
+    return middleware
+
+
+@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 span must
+    parent to the ambient current span rather than become an orphan root.
+    SDK-defined: SEP-414 only covers the traceparent-present case."""
+
+    # 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]
+    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."""
+
+    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]
+    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)
+    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):
+    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 span.attributes is not None
+    assert span.attributes["error.type"] == str(exc.value.error.code)
+    assert span.attributes["rpc.response.status_code"] == str(exc.value.error.code)
+    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)
+    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 span.attributes is not None
+    assert span.attributes["error.type"] == str(INVALID_PARAMS)
+    assert span.attributes["rpc.response.status_code"] == str(INVALID_PARAMS)
+    assert span.attributes["gen_ai.operation.name"] == "execute_tool"
+    assert "gen_ai.tool.name" not in span.attributes
+    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)
+    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"
+    assert span.attributes is not None
+    assert span.attributes["error.type"] == "ValueError"
+    [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_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] = {}
+
+    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.append(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 == "tools/call mytool"