From 675a23128c8b5930f4775c18a548eae850bdc05f Mon Sep 17 00:00:00 2001
From: zzstoatzz <thrast36@gmail.com>
Date: Thu, 9 Jul 2026 15:33:30 -0500
Subject: [PATCH] Expose sampling attributes on span start
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🤖 Generated with Codex
---
 fastmcp_slim/fastmcp/server/sampling/run.py   | 14 ++++---
 .../server/telemetry/test_sampling_tracing.py | 41 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/fastmcp_slim/fastmcp/server/sampling/run.py b/fastmcp_slim/fastmcp/server/sampling/run.py
index cb0a1e9ec..382bf8a47 100644
--- a/fastmcp_slim/fastmcp/server/sampling/run.py
+++ b/fastmcp_slim/fastmcp/server/sampling/run.py
@@ -313,10 +313,11 @@ async def execute_tools(
         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,
+            },
         ) 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(
@@ -556,12 +557,13 @@ async def sample_step_impl(
     with tracer.start_as_current_span(
         "sampling create_message",
         kind=SpanKind.CLIENT,
+        attributes={
+            "mcp.method.name": "sampling/createMessage",
+            "fastmcp.server.name": context.fastmcp.name,
+        },
         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(
diff --git a/tests/server/telemetry/test_sampling_tracing.py b/tests/server/telemetry/test_sampling_tracing.py
index db02b13bc..2a550ffb3 100644
--- a/tests/server/telemetry/test_sampling_tracing.py
+++ b/tests/server/telemetry/test_sampling_tracing.py
@@ -10,6 +10,9 @@ 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.trace import StatusCode
 
@@ -17,6 +20,28 @@ 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 {})
+
+
+@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]
 
@@ -27,7 +52,9 @@ def _exception_events(span):
 
 class TestSamplingCreateMessageSpan:
     async def test_success_creates_span_with_attributes(
-        self, trace_exporter: InMemorySpanExporter
+        self,
+        trace_exporter: InMemorySpanExporter,
+        on_start_recorder: OnStartRecorder,
     ):
         def sampling_handler(
             messages: list[SamplingMessage],
@@ -52,6 +79,10 @@ class TestSamplingCreateMessageSpan:
         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
@@ -93,7 +124,9 @@ class TestSamplingCreateMessageSpan:
 
 class TestSamplingToolSpan:
     async def test_tool_error_span_records_exception_once(
-        self, trace_exporter: InMemorySpanExporter
+        self,
+        trace_exporter: InMemorySpanExporter,
+        on_start_recorder: OnStartRecorder,
     ):
         from mcp_types import CreateMessageResultWithTools, ToolUseContent
 
@@ -146,6 +179,10 @@ class TestSamplingToolSpan:
         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.