ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-09 07:09:11 +02:00
Code Issues Projects Releases Packages Wiki Activity

Reapply span attributes after creation to survive non-forwarding samplers

Browse source 
Tracer.start_span builds the span from sampling_result.attributes, not
the attributes kwarg — a custom Sampler that returns
SamplingResult(RECORD_AND_SAMPLE) without forwarding attributes
silently drops everything FastMCP passed at creation time. Reapply the
same attributes immediately after span creation (guarded by
is_recording()) so on_start hooks and samplers still see them, while
the finished span is guaranteed to carry FastMCP's telemetry
regardless of sampler behavior.
This commit is contained in:
Jeremiah Lowin 2026-07-18 19:52:53 -04:00
commit 4207802920
No known key found for this signature in database
5 changed files with 333 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

10
fastmcp_slim/fastmcp/client/telemetry.py
View file

@ -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:

39
fastmcp_slim/fastmcp/server/sampling/run.py
View file

@ -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(

20
fastmcp_slim/fastmcp/server/telemetry.py
View file

@ -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:

141
tests/server/telemetry/test_sampling_tracing.py
View file

@ -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"] == {}

131
tests/telemetry/test_span_attributes.py
View file

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