From ed8fa0e004b208b4d727a6133bbf518e5d0486c6 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 27 Oct 2026 18:52:28 +0000 Subject: [PATCH 1/2] fix: stop native tool queues during function failure cleanup --- src/agents/run_internal/tool_execution.py | 13 + src/agents/run_internal/tool_planning.py | 145 +++++++--- src/agents/util/_asyncio_tasks.py | 20 +- tests/test_asyncio_tasks.py | 62 +++++ tests/test_native_tool_failure_history.py | 315 +++++++++++++++++++++- 5 files changed, 508 insertions(+), 47 deletions(-) diff --git a/src/agents/run_internal/tool_execution.py b/src/agents/run_internal/tool_execution.py index 7f8e86ce54..d47e98db5d 100644 --- a/src/agents/run_internal/tool_execution.py +++ b/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, diff --git a/src/agents/run_internal/tool_planning.py b/src/agents/run_internal/tool_planning.py index 4f247c92fc..e992f33d4c 100644 --- a/src/agents/run_internal/tool_planning.py +++ b/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 @@ -61,6 +62,7 @@ execute_local_shell_calls, execute_shell_calls, get_mapping_or_attr, + log_tool_action_error, ) T = TypeVar("T") @@ -984,9 +986,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 +1054,74 @@ 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, ) else: ( diff --git a/src/agents/util/_asyncio_tasks.py b/src/agents/util/_asyncio_tasks.py index f557f672c8..6d4120678e 100644 --- a/src/agents/util/_asyncio_tasks.py +++ b/src/agents/util/_asyncio_tasks.py @@ -30,6 +30,7 @@ async def gather_with_cancel( /, *, on_child_failure: Callable[[BaseException], None] | None = None, + cancelled_children: set[asyncio.Task[Any]] | None = None, ) -> tuple[T1, T2]: ... @@ -41,6 +42,7 @@ async def gather_with_cancel( /, *, on_child_failure: Callable[[BaseException], None] | None = None, + cancelled_children: set[asyncio.Task[Any]] | None = None, ) -> tuple[T1, T2, T3]: ... @@ -53,6 +55,7 @@ async def gather_with_cancel( /, *, on_child_failure: Callable[[BaseException], None] | None = None, + cancelled_children: set[asyncio.Task[Any]] | None = None, ) -> tuple[T1, T2, T3, T4]: ... @@ -66,6 +69,7 @@ async def gather_with_cancel( /, *, on_child_failure: Callable[[BaseException], None] | None = None, + cancelled_children: set[asyncio.Task[Any]] | None = None, ) -> tuple[T1, T2, T3, T4, T5]: ... @@ -80,6 +84,7 @@ async def gather_with_cancel( /, *, on_child_failure: Callable[[BaseException], None] | None = None, + cancelled_children: set[asyncio.Task[Any]] | None = None, ) -> tuple[T1, T2, T3, T4, T5, T6]: ... @@ -87,22 +92,31 @@ async def gather_with_cancel( async def gather_with_cancel( *awaitables: Awaitable[T], on_child_failure: Callable[[BaseException], None] | None = None, + cancelled_children: set[asyncio.Task[Any]] | None = None, ) -> 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, ) -> 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 drains them without cancelling twice; + cancellation of the parent always cancels all unfinished children. + """ 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 @@ -111,7 +125,9 @@ async def gather_with_cancel( raise except BaseException: for task in tasks: - if not task.done(): + if not task.done() and not ( + child_failed and cancelled_children is not None and task in cancelled_children + ): task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise diff --git a/tests/test_asyncio_tasks.py b/tests/test_asyncio_tasks.py index fb001b9c9c..26bbd82630 100644 --- a/tests/test_asyncio_tasks.py +++ b/tests/test_asyncio_tasks.py @@ -236,3 +236,65 @@ async def test_run_producer_consumer_drains_children_on_parent_cancellation() -> finally: parent.cancel() await asyncio.gather(parent, *children, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_parent", [False, True]) +async def test_gather_drains_owned_child_cancellation_but_preserves_parent_cancellation( + cancel_parent: bool, +) -> None: + started = asyncio.Event() + draining = asyncio.Event() + release = asyncio.Event() + child_failure = asyncio.Event() + interrupted = asyncio.Event() + finished = asyncio.Event() + + async def child() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + draining.set() + try: + await release.wait() + finished.set() + except asyncio.CancelledError: + interrupted.set() + raise + + async def fail() -> None: + await draining.wait() + raise RuntimeError("selected failure") + + child_task = asyncio.create_task(child()) + task = asyncio.create_task( + gather_with_cancel( + child_task, + fail(), + on_child_failure=lambda _error: child_failure.set(), + cancelled_children={child_task}, + ) + ) + try: + await asyncio.wait_for(started.wait(), 2) + child_task.cancel() + await asyncio.wait_for(child_failure.wait(), 2) + assert not interrupted.is_set() + if cancel_parent: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + assert interrupted.is_set() + else: + release.set() + with pytest.raises(RuntimeError, match="selected failure"): + await asyncio.wait_for(task, 2) + assert finished.is_set() + assert not interrupted.is_set() + finally: + release.set() + for pending in (task, child_task): + if not pending.done(): + pending.cancel() + await asyncio.gather(task, child_task, return_exceptions=True) diff --git a/tests/test_native_tool_failure_history.py b/tests/test_native_tool_failure_history.py index 07b8b7355e..cf9352e6e9 100644 --- a/tests/test_native_tool_failure_history.py +++ b/tests/test_native_tool_failure_history.py @@ -121,13 +121,13 @@ async def test_native_finalization_survives_sibling_failure( effects: list[str] = [] gather = tool_planning.gather_with_cancel - async def observe_failure(*args, on_child_failure=None): + async def observe_failure(*args, on_child_failure=None, **kwargs): def notify(error): if on_child_failure is not None: on_child_failure(error) category_failed.set() - return await gather(*args, on_child_failure=notify) + return await gather(*args, on_child_failure=notify, **kwargs) # Observe the existing category-failure boundary without replacing its cancellation. monkeypatch.setattr(tool_planning, "gather_with_cancel", observe_failure) @@ -291,13 +291,13 @@ async def test_input_verdict_during_tool_drain_preserves_selected_tool_error( effects: list[str] = [] gather = tool_planning.gather_with_cancel - async def observe_failure(*args, on_child_failure=None): + async def observe_failure(*args, on_child_failure=None, **kwargs): def notify(error): if on_child_failure is not None: on_child_failure(error) category_failed.set() - return await gather(*args, on_child_failure=notify) + return await gather(*args, on_child_failure=notify, **kwargs) monkeypatch.setattr(tool_planning, "gather_with_cancel", observe_failure) @@ -428,3 +428,310 @@ async def extract(_context): assert len(saved_outputs) == 1 finally: session.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("reject_input", [False, True]) +@pytest.mark.parametrize("cleanup_error", [False, True]) +@pytest.mark.parametrize("kind", ["custom", "local_shell"]) +async def test_selected_function_failure_stops_queued_native_invocations( + monkeypatch, streaming, reject_input, cleanup_error, kind +): + from agents import GuardrailFunctionOutput, input_guardrail + from agents.run_internal import tool_execution + + function_finalizing = asyncio.Event() + release_function = asyncio.Event() + function_finalized = asyncio.Event() + native_started = asyncio.Event() + release_native = asyncio.Event() + native_settled = asyncio.Event() + native_cancelled = asyncio.Event() + selected = asyncio.Event() + verdict_returned = asyncio.Event() + effects: list[str] = [] + original_drain = ( + tool_execution._FunctionToolBatchExecutor._raise_failure_after_draining_siblings + ) + + async def observe_selection(self, failure): + selected.set() + return await original_drain(self, failure) + + monkeypatch.setattr( + tool_execution._FunctionToolBatchExecutor, + "_raise_failure_after_draining_siblings", + observe_selection, + ) + # Event-driven assertions must finish before the production cleanup deadline. + monkeypatch.setattr(tool_execution, "_FUNCTION_TOOL_POST_INVOKE_WAIT_SECONDS", 5) + + @tool + async def completed() -> str: + return "completed output" + + @tool(failure_error_function=None) + async def fail() -> str: + await function_finalizing.wait() + await native_started.wait() + raise ValueError("selected function failure") + + class Hooks(RunHooks): + async def on_tool_end(self, context, agent, tool, result): + if tool is completed: + function_finalizing.set() + await release_function.wait() + function_finalized.set() + + async def invoke(context, value): + effects.append(value) + if value == "first": + native_started.set() + try: + await release_native.wait() + except asyncio.CancelledError: + native_cancelled.set() + if cleanup_error: + raise RuntimeError("native cleanup failed") from None + raise + return "native output" + + executor_name = "execute_custom_tool_calls" if kind == "custom" else "execute_local_shell_calls" + original_execute = getattr(tool_planning, executor_name) + + async def observe_native_completion(**kwargs): + try: + return await original_execute(**kwargs) + finally: + native_settled.set() + + monkeypatch.setattr(tool_planning, executor_name, observe_native_completion) + if kind == "custom": + native = CustomTool(name="native", description="Synthetic tool", on_invoke_tool=invoke) + native_calls = [ + ResponseCustomToolCall( + type="custom_tool_call", name="native", call_id=value, input=value + ) + for value in ("first", "queued") + ] + else: + + async def execute_shell(request): + return await invoke(request.ctx_wrapper, request.data.call_id) + + native = LocalShellTool(executor=execute_shell) + native_calls = [ + LocalShellCall( + id=value, + type="local_shell_call", + call_id=value, + status="completed", + action=LocalShellCallAction(type="exec", command=["synthetic"], env={}), + ) + for value in ("first", "queued") + ] + + @input_guardrail + async def verdict(context, agent, input): + await selected.wait() + verdict_returned.set() + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=reject_input) + + calls = [ + function_call("completed", {}, call_id="completed"), + function_call("fail", {}, call_id="failed"), + *native_calls, + ] + agent = Agent( + name="cancel-queued-native", + model=ScriptedModel([get_exact_output_stream_step(calls) if streaming else calls]), + tools=[completed, fail, native], + input_guardrails=[verdict], + ) + session = SQLiteSession("cancel-queued-native") + streamed = None + + async def run(): + nonlocal streamed + if streaming: + streamed = Runner.run_streamed(agent, "go", hooks=Hooks(), session=session) + async for _ in streamed.stream_events(): + pass + else: + await Runner.run(agent, "go", hooks=Hooks(), session=session) + + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(verdict_returned.wait(), 2) + release_native.set() + await asyncio.wait_for(native_settled.wait(), 2) + assert native_cancelled.is_set() + assert effects == ["first"] + assert not function_finalized.is_set() + release_function.set() + with pytest.raises(UserError, match="selected function failure") as caught: + await asyncio.wait_for(task, 2) + assert function_finalized.is_set() + assert caught.value.run_data is not None + outputs = [ + item for item in caught.value.run_data.new_items if item.type == "tool_call_output_item" + ] + if reject_input: + assert outputs == [] + else: + assert outputs[0].output == "completed output" + assert len(outputs) == (2 if cleanup_error and kind == "custom" else 1) + assert len(await session.get_items()) == (1 if reject_input else 1 + 2 * len(outputs)) + if streamed is not None and reject_input: + assert streamed.new_items == [] + finally: + release_native.set() + release_function.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + session.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_parent_cancellation_after_function_failure_cancels_pending_finalization( + monkeypatch, streaming +): + from agents.run_internal import tool_execution + + entered = asyncio.Event() + selected = asyncio.Event() + release = asyncio.Event() + cancelled = asyncio.Event() + original_drain = ( + tool_execution._FunctionToolBatchExecutor._raise_failure_after_draining_siblings + ) + + async def observe_selection(self, failure): + selected.set() + return await original_drain(self, failure) + + monkeypatch.setattr( + tool_execution._FunctionToolBatchExecutor, + "_raise_failure_after_draining_siblings", + observe_selection, + ) + # A parent cancellation must not wait for this sibling-finalization budget. + monkeypatch.setattr(tool_execution, "_FUNCTION_TOOL_POST_INVOKE_WAIT_SECONDS", 5) + + @tool + async def completed() -> str: + return "completed" + + @tool(failure_error_function=None) + async def fail() -> str: + await entered.wait() + raise ValueError("selected function failure") + + class Hooks(RunHooks): + async def on_tool_end(self, context, agent, tool, result): + if tool is completed: + entered.set() + try: + await release.wait() + except asyncio.CancelledError: + cancelled.set() + raise + + calls = [ + function_call("completed", {}, call_id="completed"), + function_call("fail", {}, call_id="failed"), + ] + agent = Agent(name="parent-cancel", tools=[completed, fail], model=ScriptedModel([calls])) + + async def run(): + if streaming: + result = Runner.run_streamed(agent, "go", hooks=Hooks()) + async for _ in result.stream_events(): + pass + else: + await Runner.run(agent, "go", hooks=Hooks()) + + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(selected.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + await asyncio.wait_for(cancelled.wait(), 2) + assert not release.is_set() + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_converted_native_cancellation_bounds_late_finalization(streaming): + started = asyncio.Event() + finalizing = asyncio.Event() + finalized_cancelled = asyncio.Event() + release = asyncio.Event() + effects: list[str] = [] + + async def invoke(context, value): + effects.append(value) + started.set() + try: + await release.wait() + except asyncio.CancelledError: + raise RuntimeError("native cleanup failed") from None + return "native output" + + native = CustomTool(name="native", description="Synthetic tool", on_invoke_tool=invoke) + + @tool(failure_error_function=None) + async def fail() -> str: + await started.wait() + raise ValueError("selected function failure") + + class Hooks(RunHooks): + async def on_tool_end(self, context, agent, tool, result): + if tool is native: + finalizing.set() + try: + await release.wait() + except asyncio.CancelledError: + finalized_cancelled.set() + raise + + calls = [ + ResponseCustomToolCall(type="custom_tool_call", name="native", call_id=value, input=value) + for value in ("first", "queued") + ] + calls.append(function_call("fail", {}, call_id="failed")) + agent = Agent( + name="bounded-native-finalization", + tools=[native, fail], + model=ScriptedModel([get_exact_output_stream_step(calls) if streaming else calls]), + ) + + async def run(): + if streaming: + result = Runner.run_streamed(agent, "go", hooks=Hooks()) + async for _ in result.stream_events(): + pass + else: + await Runner.run(agent, "go", hooks=Hooks()) + + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(finalizing.wait(), 2) + with pytest.raises(UserError, match="selected function failure"): + await asyncio.wait_for(task, 2) + assert finalized_cancelled.is_set() + assert effects == ["first"] + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) From 54ba6623aad77f5cef4c02a3fcfcd209a6bbdc82 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 27 Oct 2026 19:09:17 +0000 Subject: [PATCH 2/2] fix: bound cleanup of already-cancelled native tasks --- src/agents/run_internal/tool_planning.py | 3 ++ src/agents/util/_asyncio_tasks.py | 35 ++++++++++--- tests/test_asyncio_tasks.py | 1 + tests/test_native_tool_failure_history.py | 63 +++++++++++++++++++++++ 4 files changed, 96 insertions(+), 6 deletions(-) diff --git a/src/agents/run_internal/tool_planning.py b/src/agents/run_internal/tool_planning.py index e992f33d4c..7dbfb2752f 100644 --- a/src/agents/run_internal/tool_planning.py +++ b/src/agents/run_internal/tool_planning.py @@ -42,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, @@ -1122,6 +1123,8 @@ def on_category_failure(error: BaseException) -> None: ), 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: ( diff --git a/src/agents/util/_asyncio_tasks.py b/src/agents/util/_asyncio_tasks.py index 6d4120678e..b0ffb5b3ac 100644 --- a/src/agents/util/_asyncio_tasks.py +++ b/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") @@ -31,6 +31,7 @@ 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]: ... @@ -43,6 +44,7 @@ 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]: ... @@ -56,6 +58,7 @@ 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]: ... @@ -70,6 +73,7 @@ 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]: ... @@ -85,6 +89,7 @@ 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]: ... @@ -93,6 +98,7 @@ 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, ...]: ... @@ -100,12 +106,14 @@ 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. ``cancelled_children`` records children whose cancellation cleanup is already - owned by the caller. A sibling failure drains them without cancelling twice; - cancellation of the parent always cancels all unfinished children. + 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) @@ -124,10 +132,25 @@ async def gather_with_cancel( # 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() and not ( - child_failed and cancelled_children is not None and task in cancelled_children - ): + 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 diff --git a/tests/test_asyncio_tasks.py b/tests/test_asyncio_tasks.py index 26bbd82630..97673a3151 100644 --- a/tests/test_asyncio_tasks.py +++ b/tests/test_asyncio_tasks.py @@ -274,6 +274,7 @@ async def fail() -> None: fail(), on_child_failure=lambda _error: child_failure.set(), cancelled_children={child_task}, + cancelled_children_timeout=5, ) ) try: diff --git a/tests/test_native_tool_failure_history.py b/tests/test_native_tool_failure_history.py index cf9352e6e9..df60b1cf62 100644 --- a/tests/test_native_tool_failure_history.py +++ b/tests/test_native_tool_failure_history.py @@ -735,3 +735,66 @@ async def run(): if not task.done(): task.cancel() await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_pending_native_invocation_cleanup_is_bounded(streaming): + started = asyncio.Event() + cleanup_entered = asyncio.Event() + cleanup_cancelled = asyncio.Event() + release = asyncio.Event() + effects: list[str] = [] + + async def invoke(context, value): + effects.append(value) + started.set() + try: + await release.wait() + except asyncio.CancelledError: + cleanup_entered.set() + try: + await release.wait() + except asyncio.CancelledError: + cleanup_cancelled.set() + raise + return "native output" + + native = CustomTool(name="native", description="Synthetic tool", on_invoke_tool=invoke) + + @tool(failure_error_function=None) + async def fail() -> str: + await started.wait() + raise ValueError("selected function failure") + + calls = [ + ResponseCustomToolCall(type="custom_tool_call", name="native", call_id=value, input=value) + for value in ("first", "queued") + ] + calls.append(function_call("fail", {}, call_id="failed")) + agent = Agent( + name="bounded-native-cleanup", + tools=[native, fail], + model=ScriptedModel([get_exact_output_stream_step(calls) if streaming else calls]), + ) + + async def run(): + if streaming: + result = Runner.run_streamed(agent, "go") + async for _ in result.stream_events(): + pass + else: + await Runner.run(agent, "go") + + task = asyncio.create_task(run()) + try: + await asyncio.wait_for(cleanup_entered.wait(), 2) + with pytest.raises(UserError, match="selected function failure"): + await asyncio.wait_for(task, 2) + assert cleanup_cancelled.is_set() + assert effects == ["first"] + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True)