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 <noreply@anthropic.com>

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

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