fix: isolate parallel function tool failures (#2584)

Author: seratch

src/agents/agent.py

@@ -3,7 +3,6 @@
 import asyncio
 import dataclasses
 import inspect
-import json
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, cast
@@ -12,7 +11,6 @@
 from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import NotRequired, TypeAlias, TypedDict
 
-from . import _debug
 from .agent_output import AgentOutputSchemaBase
 from .agent_tool_input import (
     AgentAsToolInput,
@@ -47,12 +45,14 @@
     FunctionToolResult,
     Tool,
     ToolErrorFunction,
-    _extract_tool_argument_json_error,
+    _build_handled_function_tool_error_handler,
+    _build_wrapped_function_tool,
+    _log_function_tool_invocation,
+    _parse_function_tool_json_input,
     default_tool_error_function,
 )
 from .tool_context import ToolContext
-from .tracing import SpanError
-from .util import _error_tracing, _transforms
+from .util import _transforms
 from .util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
@@ -547,43 +547,34 @@ def _is_supported_parameters(value: Any) -> bool:
             include_json_schema=include_schema,
         )
 
-        def _normalize_tool_input(parsed: Any) -> Any:
+        def _normalize_tool_input(parsed: Any, tool_name: str) -> Any:
             # Prefer JSON mode so structured params (datetime/UUID/Decimal, etc.) serialize cleanly.
             try:
                 return params_adapter.dump_python(parsed, mode="json")
             except Exception as exc:
                 raise ModelBehaviorError(
-                    f"Failed to serialize structured tool input for {tool_name_resolved}: {exc}"
+                    f"Failed to serialize structured tool input for {tool_name}: {exc}"
                 ) from exc
 
         async def _run_agent_impl(context: ToolContext, input_json: str) -> Any:
             from .run import DEFAULT_MAX_TURNS, Runner
             from .tool_context import ToolContext
 
-            try:
-                json_data = json.loads(input_json) if input_json else {}
-            except Exception as exc:
-                if _debug.DONT_LOG_TOOL_DATA:
-                    logger.debug(f"Invalid JSON input for tool {tool_name_resolved}")
-                else:
-                    logger.debug(f"Invalid JSON input for tool {tool_name_resolved}: {input_json}")
-                raise ModelBehaviorError(
-                    f"Invalid JSON input for tool {tool_name_resolved}: {input_json}"
-                ) from exc
-
-            if _debug.DONT_LOG_TOOL_DATA:
-                logger.debug(f"Invoking tool {tool_name_resolved}")
-            else:
-                logger.debug(f"Invoking tool {tool_name_resolved} with input {input_json}")
+            tool_name = (
+                context.tool_name if isinstance(context, ToolContext) else tool_name_resolved
+            )
+            json_data = _parse_function_tool_json_input(
+                tool_name=tool_name,
+                input_json=input_json,
+            )
+            _log_function_tool_invocation(tool_name=tool_name, input_json=input_json)
 
             try:
                 parsed_params = params_adapter.validate_python(json_data)
             except ValidationError as exc:
-                raise ModelBehaviorError(
-                    f"Invalid JSON input for tool {tool_name_resolved}: {exc}"
-                ) from exc
+                raise ModelBehaviorError(f"Invalid JSON input for tool {tool_name}: {exc}") from exc
 
-            params_data = _normalize_tool_input(parsed_params)
+            params_data = _normalize_tool_input(parsed_params, tool_name)
             resolved_input = await resolve_agent_tool_input(
                 params=params_data,
                 schema_info=schema_info if should_capture_tool_input else None,
@@ -804,48 +795,17 @@ async def dispatch_stream_events() -> None:
 
             return run_result.final_output
 
-        async def _run_agent_tool(context: ToolContext, input_json: str) -> Any:
-            try:
-                return await _run_agent_impl(context, input_json)
-            except Exception as exc:
-                if failure_error_function is None:
-                    raise
-
-                result = failure_error_function(context, exc)
-                if inspect.isawaitable(result):
-                    result = await result
-
-                json_decode_error = _extract_tool_argument_json_error(exc)
-                if json_decode_error is not None:
-                    span_error_message = "Error running tool"
-                    span_error_detail = str(json_decode_error)
-                else:
-                    span_error_message = "Error running tool (non-fatal)"
-                    span_error_detail = str(exc)
-
-                _error_tracing.attach_error_to_current_span(
-                    SpanError(
-                        message=span_error_message,
-                        data={
-                            "tool_name": tool_name_resolved,
-                            "error": span_error_detail,
-                        },
-                    )
-                )
-                if _debug.DONT_LOG_TOOL_DATA:
-                    logger.debug(f"Tool {tool_name_resolved} failed")
-                else:
-                    logger.error(
-                        f"Tool {tool_name_resolved} failed: {input_json} {exc}",
-                        exc_info=exc,
-                    )
-                return result
-
-        run_agent_tool = FunctionTool(
+        run_agent_tool = _build_wrapped_function_tool(
             name=tool_name_resolved,
             description=tool_description_resolved,
             params_json_schema=params_schema,
-            on_invoke_tool=_run_agent_tool,
+            invoke_tool_impl=_run_agent_impl,
+            on_handled_error=_build_handled_function_tool_error_handler(
+                span_message="Error running tool (non-fatal)",
+                span_message_for_json_decode_error="Error running tool",
+                log_label="Tool",
+            ),
+            failure_error_function=failure_error_function,
             strict_json_schema=True,
             is_enabled=is_enabled,
             needs_approval=needs_approval,

src/agents/extensions/experimental/codex/codex_tool.py

@@ -20,11 +20,16 @@
 from agents.models import _openai_shared
 from agents.run_context import RunContextWrapper
 from agents.strict_schema import ensure_strict_json_schema
-from agents.tool import FunctionTool, ToolErrorFunction, default_tool_error_function
+from agents.tool import (
+    FunctionTool,
+    ToolErrorFunction,
+    _build_handled_function_tool_error_handler,
+    _build_wrapped_function_tool,
+    default_tool_error_function,
+)
 from agents.tool_context import ToolContext
 from agents.tracing import SpanError, custom_span
 from agents.usage import Usage as AgentsUsage
-from agents.util import _error_tracing
 from agents.util._types import MaybeAwaitable
 
 from .codex import Codex
@@ -379,7 +384,7 @@ async def _on_invoke_tool(ctx: ToolContext[Any], input_json: str) -> Any:
                     resolved_options.span_data_max_chars,
                     resolved_thread_id_holder=resolved_thread_id_holder,
                 )
-            except Exception:
+            except BaseException:
                 resolved_thread_id = resolved_thread_id_holder["thread_id"]
                 raise
 
@@ -394,38 +399,27 @@ async def _on_invoke_tool(ctx: ToolContext[Any], input_json: str) -> Any:
                 )
 
             return CodexToolResult(thread_id=resolved_thread_id, response=response, usage=usage)
-        except Exception as exc:  # noqa: BLE001
+        except BaseException:
             _try_store_thread_id_in_run_context_after_error(
                 ctx=ctx,
                 key=resolved_run_context_thread_id_key,
                 thread_id=resolved_thread_id,
                 enabled=resolved_options.use_run_context_thread_id,
             )
+            raise
 
-            if resolved_options.failure_error_function is None:
-                raise
-
-            result = resolved_options.failure_error_function(ctx, exc)
-            if inspect.isawaitable(result):
-                result = await result
-
-            _error_tracing.attach_error_to_current_span(
-                SpanError(
-                    message="Error running Codex tool (non-fatal)",
-                    data={"tool_name": name, "error": str(exc)},
-                )
-            )
-            if _debug.DONT_LOG_TOOL_DATA:
-                logger.debug("Codex tool failed")
-            else:
-                logger.error("Codex tool failed: %s", exc, exc_info=exc)
-            return result
-
-    function_tool = FunctionTool(
+    function_tool = _build_wrapped_function_tool(
         name=name,
         description=description,
         params_json_schema=params_schema,
-        on_invoke_tool=_on_invoke_tool,
+        invoke_tool_impl=_on_invoke_tool,
+        on_handled_error=_build_handled_function_tool_error_handler(
+            span_message="Error running Codex tool (non-fatal)",
+            log_label="Codex tool",
+            include_input_json_in_logs=False,
+            include_tool_name_in_log_messages=False,
+        ),
+        failure_error_function=resolved_options.failure_error_function,
         strict_json_schema=True,
         is_enabled=resolved_options.is_enabled,
     )

src/agents/mcp/util.py

@@ -22,11 +22,11 @@
     ToolErrorFunction,
     ToolOutputImageDict,
     ToolOutputTextDict,
+    _build_handled_function_tool_error_handler,
+    _build_wrapped_function_tool,
     default_tool_error_function,
 )
-from ..tool_context import ToolContext
-from ..tracing import FunctionSpanData, SpanError, get_current_span, mcp_tools_span
-from ..util import _error_tracing
+from ..tracing import FunctionSpanData, get_current_span, mcp_tools_span
 from ..util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
@@ -261,55 +261,24 @@ def to_function_tool(
             except Exception as e:
                 logger.info(f"Error converting MCP schema to strict mode: {e}")
 
-        # Wrap the invoke function with error handling, similar to regular function tools.
-        # This ensures that MCP tool errors (like timeouts) are handled gracefully instead
-        # of halting the entire agent flow.
-        async def invoke_func(ctx: ToolContext[Any], input_json: str) -> ToolOutput:
-            try:
-                return await invoke_func_impl(ctx, input_json)
-            except Exception as e:
-                if effective_failure_error_function is None:
-                    raise
-
-                # Use configured error handling function to convert exception to error message.
-                result = effective_failure_error_function(ctx, e)
-                if inspect.isawaitable(result):
-                    result = await result
-
-                # Attach error to tracing span.
-                _error_tracing.attach_error_to_current_span(
-                    SpanError(
-                        message="Error running tool (non-fatal)",
-                        data={
-                            "tool_name": tool.name,
-                            "error": str(e),
-                        },
-                    )
-                )
-
-                # Log the error.
-                if _debug.DONT_LOG_TOOL_DATA:
-                    logger.debug(f"MCP tool {tool.name} failed")
-                else:
-                    logger.error(
-                        f"MCP tool {tool.name} failed: {input_json} {e}",
-                        exc_info=e,
-                    )
-
-                return result
-
         needs_approval: (
             bool | Callable[[RunContextWrapper[Any], dict[str, Any], str], Awaitable[bool]]
         ) = server._get_needs_approval_for_tool(tool, agent)
 
-        return FunctionTool(
+        function_tool = _build_wrapped_function_tool(
             name=tool.name,
             description=tool.description or "",
             params_json_schema=schema,
-            on_invoke_tool=invoke_func,
+            invoke_tool_impl=invoke_func_impl,
+            on_handled_error=_build_handled_function_tool_error_handler(
+                span_message="Error running tool (non-fatal)",
+                log_label="MCP tool",
+            ),
+            failure_error_function=effective_failure_error_function,
             strict_json_schema=is_strict,
             needs_approval=needs_approval,
         )
+        return function_tool
 
     @staticmethod
     def _merge_mcp_meta(