fix: stop native tool queues during function failure cleanup

src/agents/run_internal/tool_execution.py

@@ -2382,6 +2382,9 @@ async def run_native_tool_post_invoke(
     task = asyncio.ensure_future(post_invoke)
     task.add_done_callback(_consume_future_exception)
     try:
+        # Invocation error conversion may consume cancellation before finalization starts.
+        if sibling_category_failure.is_set():
+            raise asyncio.CancelledError
         await asyncio.shield(task)
     except asyncio.CancelledError:
         if sibling_category_failure.is_set():
@@ -2417,6 +2420,8 @@ async def execute_custom_tool_calls(
 
     results: list[RunItem] = []
     for call in calls:
+        if sibling_category_failure is not None and sibling_category_failure.is_set():
+            break
         results.append(
             await CustomToolAction.execute(
                 agent=public_agent,
@@ -2446,6 +2451,8 @@ async def execute_local_shell_calls(
 
     results: list[RunItem] = []
     for call in calls:
+        if sibling_category_failure is not None and sibling_category_failure.is_set():
+            break
         results.append(
             await LocalShellAction.execute(
                 agent=public_agent,
@@ -2475,6 +2482,8 @@ async def execute_shell_calls(
 
     results: list[RunItem] = []
     for call in calls:
+        if sibling_category_failure is not None and sibling_category_failure.is_set():
+            break
         results.append(
             await ShellAction.execute(
                 agent=public_agent,
@@ -2504,6 +2513,8 @@ async def execute_apply_patch_calls(
 
     results: list[RunItem] = []
     for call in calls:
+        if sibling_category_failure is not None and sibling_category_failure.is_set():
+            break
         results.append(
             await ApplyPatchAction.execute(
                 agent=public_agent,
@@ -2533,6 +2544,8 @@ async def execute_computer_actions(
 
     results: list[RunItem] = []
     for action in actions:
+        if sibling_category_failure is not None and sibling_category_failure.is_set():
+            break
         context_wrapper._mark_tool_invocation_executed(
             action.tool_call,
             tool_name=action.computer_tool.name,

src/agents/run_internal/tool_planning.py

@@ -5,6 +5,7 @@
 import inspect
 import json
 from collections.abc import Awaitable, Callable, Hashable, Mapping, Sequence
+from functools import partial
 from typing import Any, TypeVar, cast
 
 from openai.types.responses import ResponseFunctionToolCall
@@ -41,6 +42,7 @@
 from ..tool import FunctionTool, MCPToolApprovalRequest, get_function_tool_origin
 from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult
 from ..util._asyncio_tasks import gather_with_cancel
+from . import tool_execution
 from .agent_bindings import AgentBindings
 from .run_steps import (
     ProcessedResponse,
@@ -61,6 +63,7 @@
     execute_local_shell_calls,
     execute_shell_calls,
     get_mapping_or_attr,
+    log_tool_action_error,
 )
 
 T = TypeVar("T")
@@ -984,9 +987,54 @@ async def _execute_tool_plan(
     )
     if parallel:
         sibling_category_failure = asyncio.Event()
+        native_category_failure = asyncio.Event()
+        native_tasks: set[asyncio.Task[Any]] = set()
+        cancelled_native_tasks: set[asyncio.Task[Any]] = set()
+        function_error_selected = False
+
+        def on_function_error_selected() -> None:
+            nonlocal function_error_selected
+            function_error_selected = True
+            native_category_failure.set()
+            for task in native_tasks:
+                if not task.done():
+                    cancelled_native_tasks.add(task)
+                    task.cancel()
+            if on_tool_error_selected is not None:
+                on_tool_error_selected()
+
+        async def run_native_category(
+            execute: Callable[[], Awaitable[list[RunItem]]],
+        ) -> list[RunItem]:
+            # A function failure may be selected before this category is scheduled.
+            if function_error_selected:
+                return []
+            task = asyncio.current_task()
+            assert task is not None
+            native_tasks.add(task)
+            try:
+                return await execute()
+            except asyncio.CancelledError:
+                if not function_error_selected:
+                    raise
+                # The function category owns this failure. Let its original error
+                # reach gather after bounded cleanup, rather than replacing it with
+                # the cancellation used to stop native invocations. Cancellation of
+                # the parent still propagates through gather_with_cancel itself.
+                return []
+            except Exception as error:
+                if task not in cancelled_native_tasks:
+                    raise
+                log_tool_action_error(
+                    "Native tool failed during sibling cancellation cleanup", error
+                )
+                return []
+            finally:
+                native_tasks.discard(task)
 
         def on_category_failure(error: BaseException) -> None:
             sibling_category_failure.set()
+            native_category_failure.set()
             if on_tool_error_selected is not None:
                 on_tool_error_selected()
             if isinstance(error, asyncio.CancelledError):
@@ -1007,58 +1055,76 @@ def on_category_failure(error: BaseException) -> None:
                 context_wrapper=context_wrapper,
                 config=run_config,
                 isolate_parallel_failures=isolate_function_tool_failures,
-                on_tool_error_selected=on_tool_error_selected,
+                on_tool_error_selected=on_function_error_selected,
                 sibling_category_failure=sibling_category_failure,
                 tool_output_committer=tool_output_committer,
                 tool_input_guardrail_results=tool_input_guardrail_results,
                 tool_output_guardrail_results=tool_output_guardrail_results,
             ),
-            execute_computer_actions(
-                public_agent=public_agent,
-                actions=plan.computer_actions,
-                hooks=hooks,
-                context_wrapper=context_wrapper,
-                config=run_config,
-                tool_output_committer=tool_output_committer,
-                sibling_category_failure=sibling_category_failure,
+            run_native_category(
+                partial(
+                    execute_computer_actions,
+                    public_agent=public_agent,
+                    actions=plan.computer_actions,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=run_config,
+                    tool_output_committer=tool_output_committer,
+                    sibling_category_failure=native_category_failure,
+                )
             ),
-            execute_custom_tool_calls(
-                public_agent=public_agent,
-                calls=plan.custom_tool_calls,
-                hooks=hooks,
-                context_wrapper=context_wrapper,
-                config=run_config,
-                tool_output_committer=tool_output_committer,
-                sibling_category_failure=sibling_category_failure,
+            run_native_category(
+                partial(
+                    execute_custom_tool_calls,
+                    public_agent=public_agent,
+                    calls=plan.custom_tool_calls,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=run_config,
+                    tool_output_committer=tool_output_committer,
+                    sibling_category_failure=native_category_failure,
+                )
             ),
-            execute_shell_calls(
-                public_agent=public_agent,
-                calls=plan.shell_calls,
-                hooks=hooks,
-                context_wrapper=context_wrapper,
-                config=run_config,
-                tool_output_committer=tool_output_committer,
-                sibling_category_failure=sibling_category_failure,
+            run_native_category(
+                partial(
+                    execute_shell_calls,
+                    public_agent=public_agent,
+                    calls=plan.shell_calls,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=run_config,
+                    tool_output_committer=tool_output_committer,
+                    sibling_category_failure=native_category_failure,
+                )
             ),
-            execute_apply_patch_calls(
-                public_agent=public_agent,
-                calls=plan.apply_patch_calls,
-                hooks=hooks,
-                context_wrapper=context_wrapper,
-                config=run_config,
-                tool_output_committer=tool_output_committer,
-                sibling_category_failure=sibling_category_failure,
+            run_native_category(
+                partial(
+                    execute_apply_patch_calls,
+                    public_agent=public_agent,
+                    calls=plan.apply_patch_calls,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=run_config,
+                    tool_output_committer=tool_output_committer,
+                    sibling_category_failure=native_category_failure,
+                )
             ),
-            execute_local_shell_calls(
-                public_agent=public_agent,
-                calls=plan.local_shell_calls,
-                hooks=hooks,
-                context_wrapper=context_wrapper,
-                config=run_config,
-                tool_output_committer=tool_output_committer,
-                sibling_category_failure=sibling_category_failure,
+            run_native_category(
+                partial(
+                    execute_local_shell_calls,
+                    public_agent=public_agent,
+                    calls=plan.local_shell_calls,
+                    hooks=hooks,
+                    context_wrapper=context_wrapper,
+                    config=run_config,
+                    tool_output_committer=tool_output_committer,
+                    sibling_category_failure=native_category_failure,
+                )
             ),
             on_child_failure=on_category_failure,
+            cancelled_children=cancelled_native_tasks,
+            # Preserve both the finalization wait and its cancellation-cleanup window.
+            cancelled_children_timeout=2 * tool_execution._FUNCTION_TOOL_POST_INVOKE_WAIT_SECONDS,
         )
     else:
         (

src/agents/util/_asyncio_tasks.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from collections.abc import Awaitable, Callable
-from typing import Any, TypeVar, overload
+from typing import Any, TypeVar, cast, overload
 
 T = TypeVar("T")
 T1 = TypeVar("T1")
@@ -30,6 +30,8 @@ async def gather_with_cancel(
     /,
     *,
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[T1, T2]: ...
 
 
@@ -41,6 +43,8 @@ async def gather_with_cancel(
     /,
     *,
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[T1, T2, T3]: ...
 
 
@@ -53,6 +57,8 @@ async def gather_with_cancel(
     /,
     *,
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[T1, T2, T3, T4]: ...
 
 
@@ -66,6 +72,8 @@ async def gather_with_cancel(
     /,
     *,
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[T1, T2, T3, T4, T5]: ...
 
 
@@ -80,38 +88,69 @@ async def gather_with_cancel(
     /,
     *,
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[T1, T2, T3, T4, T5, T6]: ...
 
 
 @overload
 async def gather_with_cancel(
     *awaitables: Awaitable[T],
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[T, ...]: ...
 
 
 async def gather_with_cancel(
     *awaitables: Awaitable[Any],
     on_child_failure: Callable[[BaseException], None] | None = None,
+    cancelled_children: set[asyncio.Task[Any]] | None = None,
+    cancelled_children_timeout: float = 0.0,
 ) -> tuple[Any, ...]:
-    """Gather awaitables, cancelling and draining siblings when one raises."""
+    """Gather awaitables, cancelling and draining siblings when one raises.
+
+    ``cancelled_children`` records children whose cancellation cleanup is already
+    owned by the caller. A sibling failure gives them ``cancelled_children_timeout``
+    seconds to finish before cancelling again. Parent cancellation always cancels
+    all unfinished children without waiting for that grace period.
+    """
     tasks = [asyncio.ensure_future(awaitable) for awaitable in awaitables]
     gather_future = asyncio.gather(*tasks)
     gather_future.add_done_callback(_consume_future_exception)
+    child_failed = False
     try:
         await asyncio.wait((gather_future,))
         try:
             return tuple(gather_future.result())
         except BaseException as error:
+            child_failed = True
             if on_child_failure is not None:
                 on_child_failure(error)
             raise
     except GeneratorExit:
         # Coroutine closure cannot suspend; the owner must tear down child tasks.
         raise
     except BaseException:
+        owned_pending: list[asyncio.Future[Any]] = []
         for task in tasks:
-            if not task.done():
+            if task.done():
+                continue
+            if child_failed and cancelled_children is not None and task in cancelled_children:
+                owned_pending.append(task)
+            else:
+                task.cancel()
+        if owned_pending:
+            try:
+                _, pending = await asyncio.wait(owned_pending, timeout=cancelled_children_timeout)
+            except BaseException:
+                # Unlike gather, wait does not forward parent cancellation to children.
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise
+            for task in cast(set[asyncio.Task[Any]], pending):
                 task.cancel()
         await asyncio.gather(*tasks, return_exceptions=True)
         raise