"""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"