mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
4207802920Tracer.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.
330 lines
12 KiB
Python
330 lines
12 KiB
Python
"""Tracing coverage for sampling create_message and tool-execution spans.
|
|
|
|
Regression focus: the `sampling create_message` span is created with
|
|
`record_exception=False, set_status_on_exception=False` and records the
|
|
exception manually in its `except` block. A failed sampling call must
|
|
therefore produce exactly ONE exception event, not two.
|
|
"""
|
|
|
|
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]
|
|
|
|
|
|
def _exception_events(span):
|
|
return [e for e in span.events if e.name == "exception"]
|
|
|
|
|
|
class TestSamplingCreateMessageSpan:
|
|
async def test_success_creates_span_with_attributes(
|
|
self,
|
|
trace_exporter: InMemorySpanExporter,
|
|
on_start_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"
|
|
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
|
|
|
|
async def test_failure_records_exception_exactly_once(
|
|
self, trace_exporter: InMemorySpanExporter
|
|
):
|
|
"""Regression: span created with record_exception=False so the manual
|
|
record_exception in the except block fires exactly once (no duplicate
|
|
exception events from OTel auto-recording on `with` exit)."""
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage],
|
|
params: SamplingParams,
|
|
ctx: RequestContext,
|
|
) -> str:
|
|
raise RuntimeError("sampling boom")
|
|
|
|
mcp = FastMCP("sampling-server")
|
|
|
|
@mcp.tool
|
|
async def ask(question: str, context: Context) -> str:
|
|
result = await context.sample(messages=question)
|
|
return result.text or ""
|
|
|
|
with pytest.raises(Exception):
|
|
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.status.status_code == StatusCode.ERROR
|
|
assert span.attributes is not None
|
|
assert "error.type" in span.attributes
|
|
# The whole point of the fix: exactly one exception event.
|
|
assert len(_exception_events(span)) == 1
|
|
|
|
|
|
class TestSamplingToolSpan:
|
|
async def test_tool_error_span_records_exception_once(
|
|
self,
|
|
trace_exporter: InMemorySpanExporter,
|
|
on_start_recorder: OnStartRecorder,
|
|
):
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
call_count = 0
|
|
|
|
def boom_tool() -> str:
|
|
raise ValueError("tool exploded")
|
|
|
|
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="boom_tool",
|
|
input={},
|
|
)
|
|
],
|
|
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=[boom_tool])
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
await client.call_tool("driver", {})
|
|
|
|
spans = _spans_named(trace_exporter, "sampling tool boom_tool")
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
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"] == {}
|