diff --git a/fastmcp_slim/fastmcp/client/telemetry.py b/fastmcp_slim/fastmcp/client/telemetry.py index d6a866098..b6fb3e464 100644 --- a/fastmcp_slim/fastmcp/client/telemetry.py +++ b/fastmcp_slim/fastmcp/client/telemetry.py @@ -42,6 +42,16 @@ def client_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(): + span.set_attributes(attrs) try: yield span except Exception as e: diff --git a/fastmcp_slim/fastmcp/server/sampling/run.py b/fastmcp_slim/fastmcp/server/sampling/run.py index 382bf8a47..0ce5524fb 100644 --- a/fastmcp_slim/fastmcp/server/sampling/run.py +++ b/fastmcp_slim/fastmcp/server/sampling/run.py @@ -310,14 +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={ - "gen_ai.tool.name": tool_use.name, - "fastmcp.tool.use_id": tool_use.id, - }, + 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_attributes(span_attrs) try: result_value = await tool.run(tool_use.input) return ToolResultContent( @@ -554,16 +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={ - "mcp.method.name": "sampling/createMessage", - "fastmcp.server.name": context.fastmcp.name, - }, + 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_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 2b994b24d..af5558f0a 100644 --- a/fastmcp_slim/fastmcp/server/telemetry.py +++ b/fastmcp_slim/fastmcp/server/telemetry.py @@ -150,6 +150,16 @@ def seam_span(method: str, server_name: str) -> Generator[Span, None, None]: 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_attributes(attrs) token = _active_seam_span.set(span) try: yield span @@ -218,6 +228,11 @@ def server_span( 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: yield span except Exception as e: @@ -246,6 +261,11 @@ def delegate_span( tracer = get_tracer() 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(): + span.set_attributes(attrs) try: yield span except Exception as e: diff --git a/tests/server/telemetry/test_sampling_tracing.py b/tests/server/telemetry/test_sampling_tracing.py index 2a550ffb3..063b8b1ee 100644 --- a/tests/server/telemetry/test_sampling_tracing.py +++ b/tests/server/telemetry/test_sampling_tracing.py @@ -14,6 +14,7 @@ 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 @@ -28,6 +29,33 @@ class OnStartRecorder(SpanProcessor): 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, @@ -187,3 +215,116 @@ class TestSamplingToolSpan: # 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 index 82fde174e..c5cbb6279 100644 --- a/tests/telemetry/test_span_attributes.py +++ b/tests/telemetry/test_span_attributes.py @@ -1,6 +1,13 @@ +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 @@ -14,6 +21,33 @@ class OnStartRecorder(SpanProcessor): 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: @@ -72,3 +106,100 @@ def test_known_span_attributes_are_available_on_start( "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