diff --git a/fastmcp_slim/fastmcp/client/telemetry.py b/fastmcp_slim/fastmcp/client/telemetry.py index e66cd7b47..b6fb3e464 100644 --- a/fastmcp_slim/fastmcp/client/telemetry.py +++ b/fastmcp_slim/fastmcp/client/telemetry.py @@ -23,23 +23,34 @@ def client_span( Automatically records any exception on the span and sets error status. """ + attrs: dict[str, str] = { + # MCP semantic conventions + "mcp.method.name": method, + # FastMCP-specific attributes + "fastmcp.component.key": component_key, + } + if session_id is not None: + attrs["mcp.session.id"] = session_id + if resource_uri: + attrs["mcp.resource.uri"] = resource_uri + if tool_name is not None: + attrs["gen_ai.tool.name"] = tool_name + if prompt_name is not None: + attrs["gen_ai.prompt.name"] = prompt_name + tracer = get_tracer() - with tracer.start_as_current_span(name, kind=SpanKind.CLIENT) as span: + with tracer.start_as_current_span( + name, kind=SpanKind.CLIENT, attributes=attrs + ) as span: + # Reapply: `attributes=attrs` above lets on_start hooks and the + # sampler see these values at creation time. But OTel's + # Tracer.start_span builds the span from + # `sampling_result.attributes`, not the `attributes` kwarg directly — + # a custom Sampler whose SamplingResult.attributes defaults to None + # silently drops everything we passed. Reapplying here (additive, + # can't clobber anything a sampler legitimately added) guarantees + # FastMCP's attributes survive regardless of sampler behavior. if span.is_recording(): - attrs: dict[str, str] = { - # MCP semantic conventions - "mcp.method.name": method, - # FastMCP-specific attributes - "fastmcp.component.key": component_key, - } - if session_id is not None: - attrs["mcp.session.id"] = session_id - if resource_uri: - attrs["mcp.resource.uri"] = resource_uri - if tool_name is not None: - attrs["gen_ai.tool.name"] = tool_name - if prompt_name is not None: - attrs["gen_ai.prompt.name"] = prompt_name span.set_attributes(attrs) try: yield span diff --git a/fastmcp_slim/fastmcp/server/sampling/run.py b/fastmcp_slim/fastmcp/server/sampling/run.py index cb0a1e9ec..0ce5524fb 100644 --- a/fastmcp_slim/fastmcp/server/sampling/run.py +++ b/fastmcp_slim/fastmcp/server/sampling/run.py @@ -310,13 +310,26 @@ async def execute_tools( ) tracer = get_tracer() + span_attrs = { + "gen_ai.tool.name": tool_use.name, + "fastmcp.tool.use_id": tool_use.id, + } with tracer.start_as_current_span( f"sampling tool {tool_use.name}", kind=SpanKind.INTERNAL, + attributes=span_attrs, ) as span: + # Reapply: `attributes=span_attrs` above lets on_start hooks and + # the sampler see these values at creation time. But OTel's + # Tracer.start_span builds the span from + # `sampling_result.attributes`, not the `attributes` kwarg + # directly — a custom Sampler whose SamplingResult.attributes + # defaults to None silently drops everything we passed. + # Reapplying here (additive, can't clobber anything a sampler + # legitimately added) guarantees FastMCP's attributes survive + # regardless of sampler behavior. if span.is_recording(): - span.set_attribute("gen_ai.tool.name", tool_use.name) - span.set_attribute("fastmcp.tool.use_id", tool_use.id) + span.set_attributes(span_attrs) try: result_value = await tool.run(tool_use.input) return ToolResultContent( @@ -553,15 +566,27 @@ async def sample_step_impl( # Make the LLM call tracer = get_tracer() + span_attrs = { + "mcp.method.name": "sampling/createMessage", + "fastmcp.server.name": context.fastmcp.name, + } with tracer.start_as_current_span( "sampling create_message", kind=SpanKind.CLIENT, + attributes=span_attrs, record_exception=False, set_status_on_exception=False, ) as span: + # Reapply: `attributes=span_attrs` above lets on_start hooks and the + # sampler see these values at creation time. But OTel's + # Tracer.start_span builds the span from + # `sampling_result.attributes`, not the `attributes` kwarg directly — + # a custom Sampler whose SamplingResult.attributes defaults to None + # silently drops everything we passed. Reapplying here (additive, + # can't clobber anything a sampler legitimately added) guarantees + # FastMCP's attributes survive regardless of sampler behavior. if span.is_recording(): - span.set_attribute("mcp.method.name", "sampling/createMessage") - span.set_attribute("fastmcp.server.name", context.fastmcp.name) + span.set_attributes(span_attrs) try: if use_fallback: response = await call_sampling_handler( diff --git a/fastmcp_slim/fastmcp/server/telemetry.py b/fastmcp_slim/fastmcp/server/telemetry.py index e9f6d4344..af5558f0a 100644 --- a/fastmcp_slim/fastmcp/server/telemetry.py +++ b/fastmcp_slim/fastmcp/server/telemetry.py @@ -135,23 +135,31 @@ def seam_span(method: str, server_name: str) -> Generator[Span, None, None]: rejections *before* the high-level path (auth, not-found, middleware vetoes) that would otherwise produce no SERVER span at all — are recorded here. """ + attrs = { + SEAM_SPAN_MARKER: True, + "mcp.method.name": method, + "fastmcp.server.name": server_name, + **get_protocol_span_attributes(), + **get_auth_span_attributes(), + **get_session_span_attributes(), + } tracer = get_tracer() with tracer.start_as_current_span( method, context=_get_parent_trace_context(), kind=SpanKind.SERVER, + attributes=attrs, ) as span: + # Reapply: `attributes=attrs` above is what makes on_start hooks and + # the sampler see these values at creation time (the whole point of + # this helper). But OTel's Tracer.start_span builds the span from + # `sampling_result.attributes`, not the `attributes` kwarg directly — + # a custom Sampler whose SamplingResult.attributes defaults to None + # silently drops everything we passed. Reapplying here (additive, + # can't clobber anything a sampler legitimately added) guarantees + # FastMCP's attributes survive regardless of sampler behavior. if span.is_recording(): - span.set_attribute(SEAM_SPAN_MARKER, True) - span.set_attributes( - { - "mcp.method.name": method, - "fastmcp.server.name": server_name, - **get_protocol_span_attributes(), - **get_auth_span_attributes(), - **get_session_span_attributes(), - } - ) + span.set_attributes(attrs) token = _active_seam_span.set(span) try: yield span @@ -218,7 +226,11 @@ def server_span( name, context=_get_parent_trace_context(), kind=SpanKind.SERVER, + attributes=attrs, ) as span: + # Reapply for the same reason as `seam_span`: OTel builds the span + # from `sampling_result.attributes`, which a custom Sampler may not + # forward even though it was handed `attributes=attrs` above. if span.is_recording(): span.set_attributes(attrs) try: @@ -240,15 +252,19 @@ def delegate_span( Used by FastMCPProvider when delegating to mounted servers. Automatically records any exception on the span and sets error status. """ + attrs: dict[str, str] = { + "fastmcp.provider.type": provider_type, + "fastmcp.component.key": component_key, + } + if method is not None: + attrs["mcp.method.name"] = method + tracer = get_tracer() - with tracer.start_as_current_span(f"delegate {name}") as span: + with tracer.start_as_current_span(f"delegate {name}", attributes=attrs) as span: + # Reapply for the same reason as `seam_span`: OTel builds the span + # from `sampling_result.attributes`, which a custom Sampler may not + # forward even though it was handed `attributes=attrs` above. if span.is_recording(): - attrs: dict[str, str] = { - "fastmcp.provider.type": provider_type, - "fastmcp.component.key": component_key, - } - if method is not None: - attrs["mcp.method.name"] = method span.set_attributes(attrs) try: yield span diff --git a/tests/server/telemetry/test_sampling_tracing.py b/tests/server/telemetry/test_sampling_tracing.py index db02b13bc..063b8b1ee 100644 --- a/tests/server/telemetry/test_sampling_tracing.py +++ b/tests/server/telemetry/test_sampling_tracing.py @@ -10,13 +10,66 @@ from __future__ import annotations import pytest from mcp_types import TextContent +from opentelemetry.context import Context as OTelContext +from opentelemetry.sdk.trace import Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult from opentelemetry.trace import StatusCode from fastmcp import Client, Context, FastMCP from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams +class OnStartRecorder(SpanProcessor): + def __init__(self) -> None: + self.attributes: dict[str, dict[str, object]] = {} + + def on_start(self, span: Span, parent_context: OTelContext | None = None) -> None: + self.attributes[span.name] = dict(span.attributes or {}) + + +class NonForwardingSampler(Sampler): + """Samples every span but never forwards the attributes it was handed. + + See `tests/telemetry/test_span_attributes.py` for the full explanation: + OTel's `Tracer.start_span` builds the finished span from + `sampling_result.attributes`, not from the `attributes` kwarg passed to + `start_as_current_span`, so a custom sampler like this one reproduces the + regression where a non-forwarding sampler silently drops FastMCP's + attributes. + """ + + def should_sample( + self, + parent_context: OTelContext | None, + trace_id: int, + name: str, + kind: object = None, + attributes: object = None, + links: object = None, + trace_state: object = None, + ) -> SamplingResult: + return SamplingResult(Decision.RECORD_AND_SAMPLE) + + def get_description(self) -> str: + return "NonForwardingSampler" + + +@pytest.fixture +def on_start_recorder( + monkeypatch: pytest.MonkeyPatch, + trace_exporter: InMemorySpanExporter, +) -> OnStartRecorder: + recorder = OnStartRecorder() + provider = TracerProvider() + provider.add_span_processor(recorder) + provider.add_span_processor(SimpleSpanProcessor(trace_exporter)) + tracer = provider.get_tracer("test") + monkeypatch.setattr("fastmcp.server.sampling.run.get_tracer", lambda: tracer) + return recorder + + def _spans_named(exporter: InMemorySpanExporter, name: str): return [s for s in exporter.get_finished_spans() if s.name == name] @@ -27,7 +80,9 @@ def _exception_events(span): class TestSamplingCreateMessageSpan: async def test_success_creates_span_with_attributes( - self, trace_exporter: InMemorySpanExporter + self, + trace_exporter: InMemorySpanExporter, + on_start_recorder: OnStartRecorder, ): def sampling_handler( messages: list[SamplingMessage], @@ -52,6 +107,10 @@ class TestSamplingCreateMessageSpan: assert span.attributes is not None assert span.attributes["mcp.method.name"] == "sampling/createMessage" assert span.attributes["fastmcp.server.name"] == "sampling-server" + assert on_start_recorder.attributes["sampling create_message"] == { + "mcp.method.name": "sampling/createMessage", + "fastmcp.server.name": "sampling-server", + } # Success path must not record any exception. assert _exception_events(span) == [] assert span.status.status_code != StatusCode.ERROR @@ -93,7 +152,9 @@ class TestSamplingCreateMessageSpan: class TestSamplingToolSpan: async def test_tool_error_span_records_exception_once( - self, trace_exporter: InMemorySpanExporter + self, + trace_exporter: InMemorySpanExporter, + on_start_recorder: OnStartRecorder, ): from mcp_types import CreateMessageResultWithTools, ToolUseContent @@ -146,7 +207,124 @@ class TestSamplingToolSpan: assert span.status.status_code == StatusCode.ERROR assert span.attributes is not None assert span.attributes["gen_ai.tool.name"] == "boom_tool" + assert on_start_recorder.attributes["sampling tool boom_tool"] == { + "gen_ai.tool.name": "boom_tool", + "fastmcp.tool.use_id": "call_1", + } assert "error.type" in span.attributes # Tool spans catch-and-convert (no re-raise), so OTel auto-recording # never fires; the manual record_exception must fire exactly once. assert len(_exception_events(span)) == 1 + + +class TestAttributesSurviveANonForwardingSampler: + """Regression: `sampling create_message` and `sampling tool ...` spans + must keep FastMCP's attributes even when the configured Sampler doesn't + forward the `attributes` it was handed to its `SamplingResult`. + """ + + @pytest.fixture + def non_forwarding_recorder( + self, + monkeypatch: pytest.MonkeyPatch, + trace_exporter: InMemorySpanExporter, + ) -> OnStartRecorder: + recorder = OnStartRecorder() + provider = TracerProvider(sampler=NonForwardingSampler()) + provider.add_span_processor(recorder) + provider.add_span_processor(SimpleSpanProcessor(trace_exporter)) + tracer = provider.get_tracer("test") + monkeypatch.setattr("fastmcp.server.sampling.run.get_tracer", lambda: tracer) + return recorder + + async def test_create_message_span_keeps_attributes( + self, + trace_exporter: InMemorySpanExporter, + non_forwarding_recorder: OnStartRecorder, + ): + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> str: + return "sampled-text" + + mcp = FastMCP("sampling-server") + + @mcp.tool + async def ask(question: str, context: Context) -> str: + result = await context.sample(messages=question) + return result.text or "" + + async with Client(mcp, sampling_handler=sampling_handler) as client: + await client.call_tool("ask", {"question": "hi"}) + + spans = _spans_named(trace_exporter, "sampling create_message") + assert len(spans) == 1 + span = spans[0] + assert span.attributes is not None + assert span.attributes["mcp.method.name"] == "sampling/createMessage" + assert span.attributes["fastmcp.server.name"] == "sampling-server" + # The sampler never forwards attributes, so on_start legitimately sees + # none — this documents that limitation rather than asserting around it. + assert non_forwarding_recorder.attributes["sampling create_message"] == {} + + async def test_sampling_tool_span_keeps_attributes( + self, + trace_exporter: InMemorySpanExporter, + non_forwarding_recorder: OnStartRecorder, + ): + from mcp_types import CreateMessageResultWithTools, ToolUseContent + + def echo_tool(text: str) -> str: + return text + + call_count = 0 + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> CreateMessageResultWithTools: + nonlocal call_count + call_count += 1 + if call_count == 1: + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="echo_tool", + input={"text": "hi"}, + ) + ], + model="test-model", + stop_reason="toolUse", + ) + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="done")], + model="test-model", + stop_reason="endTurn", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def driver(context: Context) -> str: + result = await context.sample(messages="go", tools=[echo_tool]) + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("driver", {}) + + spans = _spans_named(trace_exporter, "sampling tool echo_tool") + assert len(spans) == 1 + span = spans[0] + assert span.attributes is not None + assert span.attributes["gen_ai.tool.name"] == "echo_tool" + assert span.attributes["fastmcp.tool.use_id"] == "call_1" + # The sampler never forwards attributes, so on_start legitimately sees + # none — this documents that limitation rather than asserting around it. + assert non_forwarding_recorder.attributes["sampling tool echo_tool"] == {} diff --git a/tests/telemetry/test_span_attributes.py b/tests/telemetry/test_span_attributes.py new file mode 100644 index 000000000..c5cbb6279 --- /dev/null +++ b/tests/telemetry/test_span_attributes.py @@ -0,0 +1,205 @@ +from collections.abc import Callable +from contextlib import AbstractContextManager + +import pytest +from opentelemetry.context import Context +from opentelemetry.sdk.trace import Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult +from opentelemetry.trace import Span as APISpan + +from fastmcp.client.telemetry import client_span +from fastmcp.server.telemetry import delegate_span, seam_span, server_span + + +class OnStartRecorder(SpanProcessor): + def __init__(self) -> None: + self.attributes: dict[str, dict[str, object]] = {} + + def on_start(self, span: Span, parent_context: Context | None = None) -> None: + self.attributes[span.name] = dict(span.attributes or {}) + + +class NonForwardingSampler(Sampler): + """Samples every span but never forwards the attributes it was handed. + + Mirrors a real-world custom sampler that builds its own `SamplingResult` + without threading through the `attributes` it received — the + `attributes` parameter defaults to `None`, so `RECORD_AND_SAMPLE` with no + `attributes` argument reproduces the regression: OTel's `Tracer.start_span` + constructs the span from `sampling_result.attributes`, not from the + `attributes` kwarg passed to `start_as_current_span`. + """ + + def should_sample( + self, + parent_context: Context | None, + trace_id: int, + name: str, + kind: object = None, + attributes: object = None, + links: object = None, + trace_state: object = None, + ) -> SamplingResult: + return SamplingResult(Decision.RECORD_AND_SAMPLE) + + def get_description(self) -> str: + return "NonForwardingSampler" + + +def test_known_span_attributes_are_available_on_start( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder = OnStartRecorder() + provider = TracerProvider() + provider.add_span_processor(recorder) + tracer = provider.get_tracer("test") + monkeypatch.setattr("fastmcp.client.telemetry.get_tracer", lambda: tracer) + monkeypatch.setattr("fastmcp.server.telemetry.get_tracer", lambda: tracer) + + with client_span( + "client test", + method="tools/call", + component_key="tool:echo@", + tool_name="echo", + ): + pass + with server_span( + "server test", + method="tools/call", + server_name="test-server", + component_type="tool", + component_key="tool:echo@", + tool_name="echo", + ): + pass + with seam_span("initialize test", server_name="test-server"): + pass + with delegate_span( + "delegate test", + provider_type="LocalProvider", + component_key="tool:echo@", + method="tools/call", + ): + pass + + assert recorder.attributes["client test"] == { + "mcp.method.name": "tools/call", + "fastmcp.component.key": "tool:echo@", + "gen_ai.tool.name": "echo", + } + assert recorder.attributes["server test"] == { + "mcp.method.name": "tools/call", + "fastmcp.server.name": "test-server", + "fastmcp.component.type": "tool", + "fastmcp.component.key": "tool:echo@", + "gen_ai.tool.name": "echo", + } + assert recorder.attributes["initialize test"] == { + "fastmcp.span.seam": True, + "mcp.method.name": "initialize test", + "fastmcp.server.name": "test-server", + } + assert recorder.attributes["delegate delegate test"] == { + "fastmcp.provider.type": "LocalProvider", + "fastmcp.component.key": "tool:echo@", + "mcp.method.name": "tools/call", + } + + +SPAN_HELPER_CASES = [ + pytest.param( + lambda: client_span( + "client test", + method="tools/call", + component_key="tool:echo@", + tool_name="echo", + ), + "client test", + { + "mcp.method.name": "tools/call", + "fastmcp.component.key": "tool:echo@", + "gen_ai.tool.name": "echo", + }, + id="client_span", + ), + pytest.param( + lambda: server_span( + "server test", + method="tools/call", + server_name="test-server", + component_type="tool", + component_key="tool:echo@", + tool_name="echo", + ), + "server test", + { + "mcp.method.name": "tools/call", + "fastmcp.server.name": "test-server", + "fastmcp.component.type": "tool", + "fastmcp.component.key": "tool:echo@", + "gen_ai.tool.name": "echo", + }, + id="server_span", + ), + pytest.param( + lambda: seam_span("initialize test", server_name="test-server"), + "initialize test", + { + "fastmcp.span.seam": True, + "mcp.method.name": "initialize test", + "fastmcp.server.name": "test-server", + }, + id="seam_span", + ), + pytest.param( + lambda: delegate_span( + "delegate test", + provider_type="LocalProvider", + component_key="tool:echo@", + method="tools/call", + ), + "delegate delegate test", + { + "fastmcp.provider.type": "LocalProvider", + "fastmcp.component.key": "tool:echo@", + "mcp.method.name": "tools/call", + }, + id="delegate_span", + ), +] + + +@pytest.mark.parametrize( + ("span_factory", "span_name", "expected_attrs"), SPAN_HELPER_CASES +) +def test_attributes_survive_a_non_forwarding_sampler( + monkeypatch: pytest.MonkeyPatch, + span_factory: Callable[[], AbstractContextManager[APISpan]], + span_name: str, + expected_attrs: dict[str, object], +) -> None: + """Regression: a custom Sampler that samples a span but doesn't forward + the `attributes` it was handed must not erase FastMCP's telemetry. + + OTel's `Tracer.start_span` builds the finished span from + `sampling_result.attributes`, not from the `attributes` kwarg passed to + `start_as_current_span`. Built-in samplers forward what they're given, but + a custom sampler can legally return `SamplingResult(attributes=None)` and + silently drop everything FastMCP passed in. The span helpers must reapply + their attributes after span creation so this can't happen. + """ + exporter = InMemorySpanExporter() + provider = TracerProvider(sampler=NonForwardingSampler()) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + monkeypatch.setattr("fastmcp.client.telemetry.get_tracer", lambda: tracer) + monkeypatch.setattr("fastmcp.server.telemetry.get_tracer", lambda: tracer) + + with span_factory(): + pass + + spans = [s for s in exporter.get_finished_spans() if s.name == span_name] + assert len(spans) == 1 + assert dict(spans[0].attributes or {}) == expected_attrs