diff --git a/src/fastmcp/client/telemetry.py b/src/fastmcp/client/telemetry.py index 10d6d825f..fbb7c94e1 100644 --- a/src/fastmcp/client/telemetry.py +++ b/src/fastmcp/client/telemetry.py @@ -22,23 +22,26 @@ def client_span( """ tracer = get_tracer() with tracer.start_as_current_span(name, kind=SpanKind.CLIENT) as span: - attrs: dict[str, str] = { - # RPC semantic conventions - "rpc.system": "mcp", - "rpc.method": method, - # MCP semantic conventions - "mcp.method.name": method, - # FastMCP-specific attributes - "fastmcp.component.key": component_key, - } - if session_id: - attrs["mcp.session.id"] = session_id - if resource_uri: - attrs["mcp.resource.uri"] = resource_uri - span.set_attributes(attrs) + if span.is_recording(): + attrs: dict[str, str] = { + # RPC semantic conventions + "rpc.system": "mcp", + "rpc.method": method, + # MCP semantic conventions + "mcp.method.name": method, + # FastMCP-specific attributes + "fastmcp.component.key": component_key, + } + if session_id: + attrs["mcp.session.id"] = session_id + if resource_uri: + attrs["mcp.resource.uri"] = resource_uri + span.set_attributes(attrs) try: yield span 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)) raise diff --git a/src/fastmcp/server/sampling/run.py b/src/fastmcp/server/sampling/run.py index 0fe55ef0f..ea25088da 100644 --- a/src/fastmcp/server/sampling/run.py +++ b/src/fastmcp/server/sampling/run.py @@ -4,6 +4,7 @@ from __future__ import annotations import inspect import json +import os from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Generic, Literal, cast @@ -26,12 +27,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 @@ -41,6 +44,10 @@ from fastmcp.utilities.types import get_cached_typeadapter logger = get_logger(__name__) +_CAPTURE_CONTENT = os.environ.get( + "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "false" +).lower() == "true" + if TYPE_CHECKING: from fastmcp.server.context import Context @@ -281,49 +288,67 @@ async def execute_tools( async def _execute_single_tool(tool_use: ToolUseContent) -> ToolResultContent: """Execute a single tool and return its result.""" - tool = tool_map.get(tool_use.name) - if tool is None: - return ToolResultContent( - type="tool_result", - toolUseId=tool_use.id, - content=[ - TextContent( - type="text", - text=f"Error: Unknown tool '{tool_use.name}'", - ) - ], - isError=True, - ) + tracer = get_tracer() + with tracer.start_as_current_span( + f"sampling.execute_tool {tool_use.name}" + ) as tool_span: + if tool_span.is_recording(): + tool_span.set_attributes({ + "gen_ai.tool.name": tool_use.name, + "gen_ai.operation.name": "execute_tool", + }) - 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.exception(f"Error calling sampling tool '{tool_use.name}'") - 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, - ) + tool = tool_map.get(tool_use.name) + if tool is None: + return ToolResultContent( + type="tool_result", + toolUseId=tool_use.id, + content=[ + TextContent( + type="text", + text=f"Error: Unknown tool '{tool_use.name}'", + ) + ], + 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.exception(f"Error calling sampling tool '{tool_use.name}'") + if tool_span.is_recording(): + tool_span.set_attribute("error.type", type(e).__qualname__) + tool_span.record_exception(e) + tool_span.set_status(Status(StatusCode.ERROR, str(e))) + 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 tool_span.is_recording(): + tool_span.set_attribute("error.type", type(e).__qualname__) + tool_span.record_exception(e) + tool_span.set_status(Status(StatusCode.ERROR, str(e))) + 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( @@ -621,123 +646,197 @@ async def sample_impl( text_response_retries = 0 consecutive_validation_failures = 0 - for _iteration in range(max_iterations): - step = await sample_step_impl( - context, - messages=current_messages, - system_prompt=system_prompt, - temperature=temperature, - max_tokens=max_tokens, - model_preferences=model_preferences, - tools=sampling_tools, - tool_choice=tool_choice, - mask_error_details=mask_error_details, - tool_concurrency=tool_concurrency, - ) + tracer = get_tracer() + result_type_name = ( + result_type.__name__ if result_type and result_type is not str else "str" + ) + n_tools = len(sampling_tools) if sampling_tools else 0 - # Check for final_response tool call for structured output - had_final_response = False - if result_type is not None and result_type is not str and step.is_tool_use: - for tool_call in step.tool_calls: - if tool_call.name == "final_response": - had_final_response = True - # Validate and return the structured result - type_adapter = get_cached_typeadapter(result_type) + with tracer.start_as_current_span( + "sampling/createMessage", kind=SpanKind.INTERNAL + ) as span: + if span.is_recording(): + span.set_attributes({ + "mcp.method.name": "sampling/createMessage", + "gen_ai.request.temperature": temperature or 0.0, + "gen_ai.request.max_tokens": max_tokens or 512, + "fastmcp.sampling.tool_count": n_tools, + "fastmcp.sampling.result_type": result_type_name, + }) - # Unwrap if we wrapped primitives (non-object schemas) - input_data = tool_call.input - original_schema = compress_schema( - type_adapter.json_schema(), prune_titles=True + if _CAPTURE_CONTENT: + if system_prompt is not None: + span.add_event( + "gen_ai.system.message", + {"gen_ai.system.message": system_prompt}, ) - if ( - original_schema.get("type") != "object" - and isinstance(input_data, dict) - and "value" in input_data - ): - input_data = input_data["value"] - - try: - validated_result = type_adapter.validate_python(input_data) - text = json.dumps( - type_adapter.dump_python(validated_result, mode="json") - ) - return SamplingResult( - text=text, - result=validated_result, - history=step.history, - ) - except ValidationError as e: - consecutive_validation_failures += 1 - if consecutive_validation_failures > _MAX_VALIDATION_RETRIES: - raise RuntimeError( - f"Structured output validation failed " - f"{consecutive_validation_failures} consecutive " - f"times for type {result_type.__name__}: {e}" - ) from e - # Validation failed - add error as tool result - step.history.append( - SamplingMessage( - role="user", - content=[ - ToolResultContent( - type="tool_result", - toolUseId=tool_call.id, - content=[ - TextContent( - type="text", - text=( - f"Validation error: {e}. " - "Please try again with valid data." - ), - ) - ], - isError=True, - ) - ], - ) + prepared = prepare_messages(messages) + for msg in prepared: + content = msg.content + if isinstance(content, TextContent): + span.add_event( + f"gen_ai.{msg.role}.message", + {f"gen_ai.{msg.role}.message": content.text}, ) - # The LLM called tools but not final_response — reset validation counter - if not had_final_response: - consecutive_validation_failures = 0 + for _iteration in range(max_iterations): + with tracer.start_as_current_span( + "sampling/createMessage step" + ) as step_span: + if step_span.is_recording(): + step_span.set_attribute( + "fastmcp.sampling.iteration", _iteration + ) - # If not a tool use response, we're done - if not step.is_tool_use: - # For structured output, the LLM must use the final_response tool - if result_type is not None and result_type is not str: - text_response_retries += 1 - if text_response_retries > _MAX_TEXT_RESPONSE_RETRIES: - raise RuntimeError( - f"Expected structured output of type {result_type.__name__}, " - "but the LLM returned a text response instead of calling " - f"the final_response tool ({text_response_retries} attempts)." - ) - # Nudge the LLM to use the tool - step.history.append( - SamplingMessage( - role="user", - content=TextContent( - type="text", - text=( - "You must call the `final_response` tool to provide " - "your answer. Do not respond with text — use the tool." - ), - ), - ) + step = await sample_step_impl( + context, + messages=current_messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=max_tokens, + model_preferences=model_preferences, + tools=sampling_tools, + tool_choice=tool_choice, + mask_error_details=mask_error_details, + tool_concurrency=tool_concurrency, ) - current_messages = step.history - continue - return SamplingResult( - text=step.text, - result=cast(ResultT, step.text if step.text else ""), - history=step.history, - ) - # Continue with the updated history - current_messages = step.history + if step_span.is_recording() and hasattr(step.response, "stopReason"): + step_span.set_attribute( + "fastmcp.sampling.stop_reason", + step.response.stopReason or "unknown", + ) - # After first iteration, reset tool_choice to auto (unless structured output is required) - if result_type is None or result_type is str: - tool_choice = None + # Check for final_response tool call for structured output + had_final_response = False + if result_type is not None and result_type is not str and step.is_tool_use: + for tool_call in step.tool_calls: + if tool_call.name == "final_response": + had_final_response = True + # Validate and return the structured result + type_adapter = get_cached_typeadapter(result_type) - raise RuntimeError(f"Sampling exceeded maximum iterations ({max_iterations})") + # Unwrap if we wrapped primitives (non-object schemas) + input_data = tool_call.input + original_schema = compress_schema( + type_adapter.json_schema(), prune_titles=True + ) + if ( + original_schema.get("type") != "object" + and isinstance(input_data, dict) + and "value" in input_data + ): + input_data = input_data["value"] + + try: + validated_result = type_adapter.validate_python(input_data) + text = json.dumps( + type_adapter.dump_python( + validated_result, mode="json" + ) + ) + if span.is_recording(): + span.set_attribute( + "fastmcp.sampling.iterations", _iteration + 1 + ) + return SamplingResult( + text=text, + result=validated_result, + history=step.history, + ) + except ValidationError as e: + consecutive_validation_failures += 1 + if span.is_recording(): + span.add_event( + "sampling.validation_failure", + { + "fastmcp.sampling.consecutive_failures": consecutive_validation_failures, + }, + ) + if ( + consecutive_validation_failures + > _MAX_VALIDATION_RETRIES + ): + raise RuntimeError( + f"Structured output validation failed " + f"{consecutive_validation_failures} consecutive " + f"times for type {result_type.__name__}: {e}" + ) from e + # Validation failed - add error as tool result + step.history.append( + SamplingMessage( + role="user", + content=[ + ToolResultContent( + type="tool_result", + toolUseId=tool_call.id, + content=[ + TextContent( + type="text", + text=( + f"Validation error: {e}. " + "Please try again with valid data." + ), + ) + ], + isError=True, + ) + ], + ) + ) + + # The LLM called tools but not final_response — reset validation counter + if not had_final_response: + consecutive_validation_failures = 0 + + # If not a tool use response, we're done + if not step.is_tool_use: + # For structured output, the LLM must use the final_response tool + if result_type is not None and result_type is not str: + text_response_retries += 1 + if span.is_recording(): + span.add_event( + "sampling.text_response_retry", + { + "fastmcp.sampling.retry_count": text_response_retries, + }, + ) + if text_response_retries > _MAX_TEXT_RESPONSE_RETRIES: + raise RuntimeError( + f"Expected structured output of type {result_type.__name__}, " + "but the LLM returned a text response instead of calling " + f"the final_response tool ({text_response_retries} attempts)." + ) + # Nudge the LLM to use the tool + step.history.append( + SamplingMessage( + role="user", + content=TextContent( + type="text", + text=( + "You must call the `final_response` tool to provide " + "your answer. Do not respond with text — use the tool." + ), + ), + ) + ) + current_messages = step.history + continue + if span.is_recording(): + span.set_attribute( + "fastmcp.sampling.iterations", _iteration + 1 + ) + return SamplingResult( + text=step.text, + result=cast(ResultT, step.text if step.text else ""), + history=step.history, + ) + + # Continue with the updated history + current_messages = step.history + + # After first iteration, reset tool_choice to auto (unless structured output is required) + if result_type is None or result_type is str: + tool_choice = None + + raise RuntimeError(f"Sampling exceeded maximum iterations ({max_iterations})") diff --git a/src/fastmcp/server/telemetry.py b/src/fastmcp/server/telemetry.py index 6c263225d..0e125956b 100644 --- a/src/fastmcp/server/telemetry.py +++ b/src/fastmcp/server/telemetry.py @@ -71,26 +71,29 @@ def server_span( context=_get_parent_trace_context(), kind=SpanKind.SERVER, ) as span: - attrs: dict[str, str] = { - # RPC semantic conventions - "rpc.system": "mcp", - "rpc.service": server_name, - "rpc.method": method, - # MCP semantic conventions - "mcp.method.name": method, - # FastMCP-specific attributes - "fastmcp.server.name": server_name, - "fastmcp.component.type": component_type, - "fastmcp.component.key": component_key, - **get_auth_span_attributes(), - **get_session_span_attributes(), - } - if resource_uri is not None: - attrs["mcp.resource.uri"] = resource_uri - span.set_attributes(attrs) + if span.is_recording(): + attrs: dict[str, str] = { + # RPC semantic conventions + "rpc.system": "mcp", + "rpc.service": server_name, + "rpc.method": method, + # MCP semantic conventions + "mcp.method.name": method, + # FastMCP-specific attributes + "fastmcp.server.name": server_name, + "fastmcp.component.type": component_type, + "fastmcp.component.key": component_key, + **get_auth_span_attributes(), + **get_session_span_attributes(), + } + if resource_uri is not None: + attrs["mcp.resource.uri"] = resource_uri + span.set_attributes(attrs) try: yield span 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)) raise @@ -109,15 +112,18 @@ def delegate_span( """ tracer = get_tracer() with tracer.start_as_current_span(f"delegate {name}") as span: - span.set_attributes( - { - "fastmcp.provider.type": provider_type, - "fastmcp.component.key": component_key, - } - ) + if span.is_recording(): + span.set_attributes( + { + "fastmcp.provider.type": provider_type, + "fastmcp.component.key": component_key, + } + ) try: yield span 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)) raise diff --git a/tests/server/telemetry/test_sampling_tracing.py b/tests/server/telemetry/test_sampling_tracing.py new file mode 100644 index 000000000..2eb872540 --- /dev/null +++ b/tests/server/telemetry/test_sampling_tracing.py @@ -0,0 +1,475 @@ +"""Tests for OpenTelemetry instrumentation of sampling/createMessage.""" + +from __future__ import annotations + +from mcp.types import ( + CreateMessageResultWithTools, + TextContent, + ToolUseContent, +) +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode +from pydantic import BaseModel + +from fastmcp import Client, Context, FastMCP +from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams + + +class TestSamplingCreateMessageSpan: + """Verify the top-level sampling/createMessage span.""" + + async def test_span_created_with_correct_attributes( + self, trace_exporter: InMemorySpanExporter + ): + """sampling/createMessage span has expected attributes.""" + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> str: + return "Hello from LLM" + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample( + "Hi", + temperature=0.7, + max_tokens=256, + ) + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + sampling_spans = [s for s in spans if s.name == "sampling/createMessage"] + assert len(sampling_spans) == 1 + + span = sampling_spans[0] + assert span.kind == SpanKind.INTERNAL + assert span.attributes is not None + assert span.attributes["mcp.method.name"] == "sampling/createMessage" + assert span.attributes["gen_ai.request.temperature"] == 0.7 + assert span.attributes["gen_ai.request.max_tokens"] == 256 + assert span.attributes["fastmcp.sampling.tool_count"] == 0 + assert span.attributes["fastmcp.sampling.result_type"] == "str" + + async def test_span_records_iteration_count( + self, trace_exporter: InMemorySpanExporter + ): + """sampling/createMessage span has fastmcp.sampling.iterations on completion.""" + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> str: + return "Done" + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("Hi") + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + sampling_spans = [s for s in spans if s.name == "sampling/createMessage"] + assert len(sampling_spans) == 1 + span = sampling_spans[0] + assert span.attributes is not None + assert span.attributes["fastmcp.sampling.iterations"] == 1 + + async def test_result_type_attribute(self, trace_exporter: InMemorySpanExporter): + """result_type is recorded on the span.""" + + class MyResult(BaseModel): + value: int + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> CreateMessageResultWithTools: + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="final_response", + input={"value": 42}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("Hi", result_type=MyResult) + return str(result.result.value) + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + sampling_spans = [s for s in spans if s.name == "sampling/createMessage"] + assert len(sampling_spans) == 1 + span = sampling_spans[0] + assert span.attributes is not None + assert span.attributes["fastmcp.sampling.result_type"] == "MyResult" + + +class TestSamplingStepSpans: + """Verify child sampling/createMessage step spans.""" + + async def test_step_span_created_per_iteration( + self, trace_exporter: InMemorySpanExporter + ): + """Each loop iteration creates a sampling/createMessage step span.""" + 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="my_tool", + input={"x": 1}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="Done")], + model="test-model", + stopReason="endTurn", + ) + + def my_tool(x: int) -> int: + return x * 2 + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("Hi", tools=[my_tool]) + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + step_spans = [s for s in spans if s.name == "sampling/createMessage step"] + assert len(step_spans) == 2 + + # First step should have iteration=0 + assert step_spans[0].attributes is not None + assert step_spans[0].attributes["fastmcp.sampling.iteration"] == 0 + + # Second step should have iteration=1 + assert step_spans[1].attributes is not None + assert step_spans[1].attributes["fastmcp.sampling.iteration"] == 1 + + async def test_step_span_records_stop_reason( + self, trace_exporter: InMemorySpanExporter + ): + """Step span records stop reason from the response.""" + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> CreateMessageResultWithTools: + 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 do_sample(context: Context) -> str: + result = await context.sample("Hi") + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + step_spans = [s for s in spans if s.name == "sampling/createMessage step"] + assert len(step_spans) == 1 + assert step_spans[0].attributes is not None + assert step_spans[0].attributes["fastmcp.sampling.stop_reason"] == "endTurn" + + +class TestSamplingToolExecutionSpans: + """Verify sampling.execute_tool spans.""" + + async def test_tool_execution_span(self, trace_exporter: InMemorySpanExporter): + """Tool call within sampling creates an execute_tool span.""" + 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="get_weather", + input={"city": "Seattle"}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="Done")], + model="test-model", + stopReason="endTurn", + ) + + def get_weather(city: str) -> str: + return f"Sunny in {city}" + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("weather?", tools=[get_weather]) + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + tool_spans = [s for s in spans if s.name == "sampling.execute_tool get_weather"] + assert len(tool_spans) == 1 + span = tool_spans[0] + assert span.attributes is not None + assert span.attributes["gen_ai.tool.name"] == "get_weather" + assert span.attributes["gen_ai.operation.name"] == "execute_tool" + + async def test_tool_execution_error_sets_status( + self, trace_exporter: InMemorySpanExporter + ): + """Tool error within sampling sets error status on the span.""" + 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="failing_tool", + input={}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="Handled")], + model="test-model", + stopReason="endTurn", + ) + + def failing_tool() -> str: + raise ValueError("boom") + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("test", tools=[failing_tool]) + return result.text or "" + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + tool_spans = [ + s for s in spans if s.name == "sampling.execute_tool failing_tool" + ] + assert len(tool_spans) == 1 + span = tool_spans[0] + assert span.status.status_code == StatusCode.ERROR + assert span.attributes is not None + assert span.attributes["error.type"] == "ValueError" + + +class TestSamplingValidationEvents: + """Verify validation failure and text response retry events.""" + + async def test_validation_failure_event(self, trace_exporter: InMemorySpanExporter): + """Validation failure adds a sampling.validation_failure event.""" + + class StrictModel(BaseModel): + count: int + + call_count = 0 + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> CreateMessageResultWithTools: + nonlocal call_count + call_count += 1 + if call_count == 1: + # First attempt: invalid data + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="final_response", + input={"count": "not_a_number"}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + # Second attempt: valid data + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_2", + name="final_response", + input={"count": 42}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("give me a count", result_type=StrictModel) + return str(result.result.count) + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + sampling_spans = [s for s in spans if s.name == "sampling/createMessage"] + assert len(sampling_spans) == 1 + span = sampling_spans[0] + + validation_events = [ + e for e in span.events if e.name == "sampling.validation_failure" + ] + assert len(validation_events) == 1 + assert validation_events[0].attributes is not None + assert ( + validation_events[0].attributes["fastmcp.sampling.consecutive_failures"] + == 1 + ) + + async def test_text_response_retry_event( + self, trace_exporter: InMemorySpanExporter + ): + """Text response instead of tool call adds a sampling.text_response_retry event.""" + + class MyResult(BaseModel): + value: int + + call_count = 0 + + def sampling_handler( + messages: list[SamplingMessage], + params: SamplingParams, + ctx: RequestContext, + ) -> CreateMessageResultWithTools: + nonlocal call_count + call_count += 1 + if call_count == 1: + # First attempt: text response instead of tool call + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="I think 42")], + model="test-model", + stopReason="endTurn", + ) + # Second attempt: correct tool call + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="final_response", + input={"value": 42}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def do_sample(context: Context) -> str: + result = await context.sample("give me a value", result_type=MyResult) + return str(result.result.value) + + async with Client(mcp) as client: + await client.call_tool("do_sample", {}) + + spans = trace_exporter.get_finished_spans() + sampling_spans = [s for s in spans if s.name == "sampling/createMessage"] + assert len(sampling_spans) == 1 + span = sampling_spans[0] + + retry_events = [ + e for e in span.events if e.name == "sampling.text_response_retry" + ] + assert len(retry_events) == 1 + assert retry_events[0].attributes is not None + assert retry_events[0].attributes["fastmcp.sampling.retry_count"] == 1