fix: #2729 avoid eager-task race in function tool batch executor (#2731)

Author: seratch

src/agents/run_internal/tool_execution.py

@@ -1334,10 +1334,15 @@ async def execute(
         )
 
     def _create_tool_task(self, tool_run: ToolRunFunction, order: int) -> None:
+        task_state = _FunctionToolTaskState(tool_run=tool_run, order=order)
         task = asyncio.create_task(
-            self._run_single_tool(tool_run.function_tool, tool_run.tool_call)
+            self._run_single_tool(
+                task_state=task_state,
+                func_tool=tool_run.function_tool,
+                tool_call=tool_run.tool_call,
+            )
         )
-        self.task_states[task] = _FunctionToolTaskState(tool_run=tool_run, order=order)
+        self.task_states[task] = task_state
         self.pending_tasks.add(task)
 
     async def _drain_pending_tasks(self) -> None:
@@ -1431,13 +1436,14 @@ def _cancel_pending_tasks_for_parent_cancellation(self) -> None:
 
     async def _run_single_tool(
         self,
+        *,
+        task_state: _FunctionToolTaskState,
         func_tool: FunctionTool,
         tool_call: ResponseFunctionToolCall,
     ) -> Any:
         raw_tool_call = tool_call
-        current_task = asyncio.current_task()
-        if current_task is not None:
-            self.task_states[current_task].in_post_invoke_phase = False
+        outer_task = asyncio.current_task()
+        task_state.in_post_invoke_phase = False
 
         tool_call = cast(
             ResponseFunctionToolCall,
@@ -1475,7 +1481,8 @@ async def _run_single_tool(
                     result = approval_result
                 else:
                     result = await self._execute_single_tool_body(
-                        current_task=current_task,
+                        outer_task=outer_task,
+                        task_state=task_state,
                         func_tool=func_tool,
                         tool_call=tool_call,
                         tool_context=tool_context,
@@ -1576,7 +1583,8 @@ async def _maybe_execute_tool_approval(
     async def _execute_single_tool_body(
         self,
         *,
-        current_task: asyncio.Task[Any] | None,
+        outer_task: asyncio.Task[Any] | None,
+        task_state: _FunctionToolTaskState,
         func_tool: FunctionTool,
         tool_call: ResponseFunctionToolCall,
         tool_context: ToolContext[Any],
@@ -1602,21 +1610,22 @@ async def _execute_single_tool_body(
 
         invoke_task = asyncio.create_task(
             self._invoke_tool_and_run_post_invoke(
-                current_task=current_task,
+                outer_task=outer_task,
+                task_state=task_state,
                 func_tool=func_tool,
                 tool_call=tool_call,
                 tool_context=tool_context,
                 agent_hooks=agent_hooks,
             )
         )
-        if current_task is not None:
-            self.task_states[current_task].invoke_task = invoke_task
-        return await self._await_invoke_task(current_task=current_task, invoke_task=invoke_task)
+        task_state.invoke_task = invoke_task
+        return await self._await_invoke_task(outer_task=outer_task, invoke_task=invoke_task)
 
     async def _invoke_tool_and_run_post_invoke(
         self,
         *,
-        current_task: asyncio.Task[Any] | None,
+        outer_task: asyncio.Task[Any] | None,
+        task_state: _FunctionToolTaskState,
         func_tool: FunctionTool,
         tool_call: ResponseFunctionToolCall,
         tool_context: ToolContext[Any],
@@ -1629,7 +1638,7 @@ async def _invoke_tool_and_run_post_invoke(
                 arguments=tool_call.arguments,
             )
         except asyncio.CancelledError as e:
-            if not self.isolate_parallel_failures or current_task in self.teardown_cancelled_tasks:
+            if not self.isolate_parallel_failures or outer_task in self.teardown_cancelled_tasks:
                 raise
 
             result = await maybe_invoke_function_tool_failure_error_function(
@@ -1648,8 +1657,7 @@ async def _invoke_tool_and_run_post_invoke(
             )
             real_result = result
 
-        if current_task is not None:
-            self.task_states[current_task].in_post_invoke_phase = True
+        task_state.in_post_invoke_phase = True
 
         final_result = await _execute_tool_output_guardrails(
             func_tool=func_tool,
@@ -1672,14 +1680,14 @@ async def _invoke_tool_and_run_post_invoke(
     async def _await_invoke_task(
         self,
         *,
-        current_task: asyncio.Task[Any] | None,
+        outer_task: asyncio.Task[Any] | None,
         invoke_task: asyncio.Task[Any],
     ) -> Any:
         try:
             return await asyncio.shield(invoke_task)
         except asyncio.CancelledError as cancel_exc:
             sibling_failure_cancelled = (
-                current_task is not None and current_task in self.teardown_cancelled_tasks
+                outer_task is not None and outer_task in self.teardown_cancelled_tasks
             )
             if not invoke_task.done():
                 invoke_task.cancel()

tests/test_run_step_execution.py

@@ -1163,6 +1163,61 @@ def _failure_handler(_ctx: RunContextWrapper[Any], error: Exception) -> str:
     assert not on_tool_end_called.is_set()
 
 
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not hasattr(asyncio, "eager_task_factory"),
+    reason="eager_task_factory requires Python 3.12+",
+)
+async def test_execute_function_tool_calls_eager_task_factory_tracks_state_safely():
+    async def _first_tool() -> str:
+        return "first"
+
+    async def _second_tool() -> str:
+        return "second"
+
+    first_tool = function_tool(_first_tool, name_override="first_tool")
+    second_tool = function_tool(_second_tool, name_override="second_tool")
+    tool_runs = [
+        ToolRunFunction(
+            tool_call=cast(
+                ResponseFunctionToolCall,
+                get_function_tool_call("first_tool", "{}", call_id="call-1"),
+            ),
+            function_tool=first_tool,
+        ),
+        ToolRunFunction(
+            tool_call=cast(
+                ResponseFunctionToolCall,
+                get_function_tool_call("second_tool", "{}", call_id="call-2"),
+            ),
+            function_tool=second_tool,
+        ),
+    ]
+    loop = asyncio.get_running_loop()
+    previous_task_factory = loop.get_task_factory()
+    eager_task_factory = cast(Any, asyncio.eager_task_factory)
+    loop.set_task_factory(eager_task_factory)
+
+    try:
+        (
+            function_results,
+            input_guardrail_results,
+            output_guardrail_results,
+        ) = await execute_function_tool_calls(
+            agent=Agent(name="test", tools=[first_tool, second_tool]),
+            tool_runs=tool_runs,
+            hooks=RunHooks(),
+            context_wrapper=RunContextWrapper(None),
+            config=RunConfig(),
+        )
+    finally:
+        loop.set_task_factory(previous_task_factory)
+
+    assert [result.output for result in function_results] == ["first", "second"]
+    assert input_guardrail_results == []
+    assert output_guardrail_results == []
+
+
 @pytest.mark.asyncio
 async def test_execute_function_tool_calls_collapse_trace_name_for_top_level_deferred_tools():
     async def _shipping_eta(tracking_number: str) -> str: