From e8a6b8f223a13d01ee3920acc04f903200f07772 Mon Sep 17 00:00:00 2001 From: Ethan Leifer Date: Mon, 19 Oct 2026 05:50:47 -0400 Subject: [PATCH] fix(pydantic-ai): don't log deferred tool calls as tool span errors pydantic_ai tools raise CallDeferred or ApprovalRequired to hand a call back to the caller for external execution or approval. The tool span recorded these exceptions as errors, so every deferred call looked like a failed tool in traces. Log these deferrals as the span output and re-raise them after the span closes. Other exceptions are still recorded as errors. Co-Authored-By: Claude Opus 5.5 --- .../test_pydantic_ai_integration.py | 62 +++++++++++++++++++ .../integrations/pydantic_ai/tracing.py | 30 ++++++++- 2 files changed, 89 insertions(+), 3 deletions(-) diff --git a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py index 8cc7d49b..727059b9 100644 --- a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py +++ b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py @@ -1617,6 +1617,68 @@ def calculate(operation: str, a: float, b: float) -> float: ) +def _call_tool_once_model(tool_name): + """FunctionModel that calls ``tool_name`` once, then answers with text. + + Deferral and tool-error behavior is raised by the local tool function, not + shaped by a provider response, so these tests drive the agent with + pydantic_ai's FunctionModel instead of a cassette. + """ + from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart + from pydantic_ai.models.function import FunctionModel + + def respond(messages, info): + if len(messages) == 1: + return ModelResponse(parts=[ToolCallPart(tool_name=tool_name, args={"flow_id": "abc"})]) + return ModelResponse(parts=[TextPart("done")]) + + return FunctionModel(respond) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exception_name", ["CallDeferred", "ApprovalRequired"]) +async def test_deferred_tool_call_is_not_logged_as_error(memory_logger, exception_name): + """CallDeferred and ApprovalRequired are control flow, not tool failures.""" + from pydantic_ai import DeferredToolRequests, exceptions + + assert not memory_logger.pop() + + agent = Agent(_call_tool_once_model("publish_flow"), output_type=[str, DeferredToolRequests]) + + @agent.tool_plain + def publish_flow(flow_id: str) -> str: + raise getattr(exceptions, exception_name)() + + result = await agent.run("Publish flow abc") + + assert isinstance(result.output, DeferredToolRequests) + + spans = memory_logger.pop() + tool_span = next(s for s in spans if s["span_attributes"]["name"] == "publish_flow") + assert tool_span["span_attributes"]["type"] == SpanTypeAttribute.TOOL + assert tool_span.get("error") is None + assert tool_span["output"] == {"deferred": True, "reason": exception_name} + assert tool_span["metrics"]["duration"] >= 0 + + +@pytest.mark.asyncio +async def test_tool_exception_is_logged_as_error(memory_logger): + assert not memory_logger.pop() + + agent = Agent(_call_tool_once_model("publish_flow")) + + @agent.tool_plain + def publish_flow(flow_id: str) -> str: + raise RuntimeError("publish failed") + + with pytest.raises(RuntimeError, match="publish failed"): + await agent.run("Publish flow abc") + + spans = memory_logger.pop() + tool_span = next(s for s in spans if s["span_attributes"]["name"] == "publish_flow") + assert "publish failed" in str(tool_span.get("error")) + + def test_agent_tool_metadata_extraction(memory_logger): """Test that agent tools are properly extracted with full schemas in INPUT (not metadata). diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 3c7a1bc0..a0586d70 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -1,5 +1,6 @@ import asyncio import contextvars +import functools import inspect import logging import sys @@ -985,15 +986,38 @@ async def _trace_tool_execution(wrapped: Any, args: Any, kwargs: Any): metadata = {"tool_call_id": tool_call_id} if tool_call_id else None _mark_tool_span_emitted() + deferral = None with start_span(name=tool_name, type=SpanTypeAttribute.TOOL, input=input_data, metadata=metadata) as tool_span: start_time = time.time() - result = await wrapped(*args, **kwargs) + try: + result = await wrapped(*args, **kwargs) + output = result + except _deferred_tool_exceptions() as exc: + # The tool handed the call back to the caller, so log the deferral + # as output and re-raise after the span closes without an error. + deferral = exc + output = {"deferred": True, "reason": type(exc).__name__} end_time = time.time() tool_span.log( - output=result, + output=output, metrics={"start": start_time, "end": end_time, "duration": end_time - start_time}, ) - return result + if deferral is not None: + raise deferral + return result + + +@functools.cache +def _deferred_tool_exceptions() -> tuple[type[Exception], ...]: + """Exceptions pydantic_ai tools raise to defer a call rather than to fail.""" + try: + from pydantic_ai.exceptions import ( # pylint: disable=import-outside-toplevel + ApprovalRequired, + CallDeferred, + ) + except ImportError: + return () + return (CallDeferred, ApprovalRequired) async def _tool_manager_call_function_tool_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):