From 2d61d8a46be5293a0d23925af755607effdd3a82 Mon Sep 17 00:00:00 2001 From: Bill Easton Date: Wed, 20 May 2026 08:38:58 -0500 Subject: [PATCH] fix: add OTEL spans to sampling step and tool execution (#4059) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 --- fastmcp_slim/fastmcp/server/sampling/run.py | 143 +++++++++------- .../server/telemetry/test_sampling_tracing.py | 152 ++++++++++++++++++ 2 files changed, 240 insertions(+), 55 deletions(-) create mode 100644 tests/server/telemetry/test_sampling_tracing.py diff --git a/fastmcp_slim/fastmcp/server/sampling/run.py b/fastmcp_slim/fastmcp/server/sampling/run.py index 385c2e315..47c04d78c 100644 --- a/fastmcp_slim/fastmcp/server/sampling/run.py +++ b/fastmcp_slim/fastmcp/server/sampling/run.py @@ -26,12 +26,14 @@ from mcp.types import ( ) from mcp.types import CreateMessageRequestParams as SamplingParams from mcp.types import Tool as SDKTool +from opentelemetry.trace import SpanKind, Status, StatusCode from pydantic import ValidationError from typing_extensions import TypeVar from fastmcp import settings from fastmcp.exceptions import ToolError from fastmcp.server.sampling.sampling_tool import SamplingTool +from fastmcp.telemetry import get_tracer from fastmcp.tools.function_tool import FunctionTool from fastmcp.tools.tool_transform import TransformedTool from fastmcp.utilities.async_utils import gather @@ -295,39 +297,53 @@ async def execute_tools( isError=True, ) - try: - result_value = await tool.run(tool_use.input) - return ToolResultContent( - type="tool_result", - toolUseId=tool_use.id, - content=[TextContent(type="text", text=str(result_value))], - ) - except ToolError as e: - # ToolError is the escape hatch - always pass message through - logger.log( - e.log_level, - f"Error calling sampling tool '{tool_use.name}'", - exc_info=True, - ) - return ToolResultContent( - type="tool_result", - toolUseId=tool_use.id, - content=[TextContent(type="text", text=str(e))], - isError=True, - ) - except Exception as e: - # Generic exceptions - mask based on setting - logger.exception(f"Error calling sampling tool '{tool_use.name}'") - if mask_error_details: - error_text = f"Error executing tool '{tool_use.name}'" - else: - error_text = f"Error executing tool '{tool_use.name}': {e}" - return ToolResultContent( - type="tool_result", - toolUseId=tool_use.id, - content=[TextContent(type="text", text=error_text)], - isError=True, - ) + tracer = get_tracer() + with tracer.start_as_current_span( + f"sampling tool {tool_use.name}", + kind=SpanKind.INTERNAL, + ) as span: + if span.is_recording(): + span.set_attribute("gen_ai.tool.name", tool_use.name) + span.set_attribute("fastmcp.tool.use_id", tool_use.id) + try: + result_value = await tool.run(tool_use.input) + return ToolResultContent( + type="tool_result", + toolUseId=tool_use.id, + content=[TextContent(type="text", text=str(result_value))], + ) + except ToolError as e: + if span.is_recording(): + span.set_attribute("error.type", "tool_error") + span.record_exception(e) + span.set_status(Status(StatusCode.ERROR, str(e))) + logger.log( + e.log_level, + f"Error calling sampling tool '{tool_use.name}'", + exc_info=True, + ) + return ToolResultContent( + type="tool_result", + toolUseId=tool_use.id, + content=[TextContent(type="text", text=str(e))], + isError=True, + ) + except Exception as e: + if span.is_recording(): + span.set_attribute("error.type", type(e).__qualname__) + span.record_exception(e) + span.set_status(Status(StatusCode.ERROR, str(e))) + logger.exception(f"Error calling sampling tool '{tool_use.name}'") + if mask_error_details: + error_text = f"Error executing tool '{tool_use.name}'" + else: + error_text = f"Error executing tool '{tool_use.name}': {e}" + return ToolResultContent( + type="tool_result", + toolUseId=tool_use.id, + content=[TextContent(type="text", text=error_text)], + isError=True, + ) # Check if any tool requires sequential execution requires_sequential = any( @@ -516,28 +532,45 @@ async def sample_step_impl( effective_max_tokens = max_tokens if max_tokens is not None else 512 # Make the LLM call - if use_fallback: - response = await call_sampling_handler( - context, - current_messages, - system_prompt=system_prompt, - temperature=temperature, - max_tokens=effective_max_tokens, - model_preferences=model_preferences, - sdk_tools=sdk_tools, - tool_choice=effective_tool_choice, - ) - else: - response = await context.session.create_message( - messages=current_messages, - system_prompt=system_prompt, - temperature=temperature, - max_tokens=effective_max_tokens, - model_preferences=_parse_model_preferences(model_preferences), - tools=sdk_tools, - tool_choice=effective_tool_choice, - related_request_id=context.origin_request_id, - ) + tracer = get_tracer() + with tracer.start_as_current_span( + "sampling create_message", + kind=SpanKind.CLIENT, + record_exception=False, + set_status_on_exception=False, + ) as span: + if span.is_recording(): + span.set_attribute("mcp.method.name", "sampling/createMessage") + span.set_attribute("fastmcp.server.name", context.fastmcp.name) + try: + if use_fallback: + response = await call_sampling_handler( + context, + current_messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=effective_max_tokens, + model_preferences=model_preferences, + sdk_tools=sdk_tools, + tool_choice=effective_tool_choice, + ) + else: + response = await context.session.create_message( + messages=current_messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=effective_max_tokens, + model_preferences=_parse_model_preferences(model_preferences), + tools=sdk_tools, + tool_choice=effective_tool_choice, + related_request_id=context.origin_request_id, + ) + except Exception as e: + if span.is_recording(): + span.set_attribute("error.type", type(e).__qualname__) + span.record_exception(e) + span.set_status(Status(StatusCode.ERROR, str(e))) + raise # Check if this is a tool use response is_tool_use_response = ( diff --git a/tests/server/telemetry/test_sampling_tracing.py b/tests/server/telemetry/test_sampling_tracing.py new file mode 100644 index 000000000..692de6f30 --- /dev/null +++ b/tests/server/telemetry/test_sampling_tracing.py @@ -0,0 +1,152 @@ +"""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.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + +from fastmcp import Client, Context, FastMCP +from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams + + +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 + ): + 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" + # 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 + ): + 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", + stopReason="toolUse", + ) + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="done")], + model="test-model", + stopReason="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 "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