fix: cancel sibling tasks on any exception in parallel tool execution (#4502)

Author: LakshmiSravyaVedantham

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -1281,7 +1281,14 @@ async def handle_call_or_result(
             except asyncio.CancelledError as e:
                 for task in tasks:
                     task.cancel(msg=e.args[0] if len(e.args) != 0 else None)
-
+                raise
+            except BaseException:
+                # Cancel any still-running sibling tasks so they don't become
+                # orphaned asyncio tasks when a non-CancelledError exception
+                # (e.g. RuntimeError, ConnectionError) propagates out of
+                # handle_call_or_result().
+                for task in tasks:
+                    task.cancel()
                 raise
 
     # We append the results at the end, rather than as they are received, to retain a consistent ordering

tests/test_agent.py

@@ -6499,6 +6499,58 @@ async def call_tools_parallel(messages: list[ModelMessage], info: AgentInfo) ->
     assert result.output == snapshot('finished')
 
 
+async def test_parallel_tool_exception_cancels_sibling_tasks():
+    """Non-CancelledError exceptions during parallel tool execution must cancel sibling tasks.
+
+    Regression test for https://github.com/pydantic/pydantic-ai/issues/4423.
+    Previously only asyncio.CancelledError triggered cleanup; any other exception
+    left the remaining tasks running as orphaned asyncio tasks.
+    """
+    slow_tool_started = asyncio.Event()
+    slow_tool_cancelled = asyncio.Event()
+
+    async def call_two_tools(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        return ModelResponse(
+            parts=[
+                ToolCallPart(tool_name='fast_failing_tool'),
+                ToolCallPart(tool_name='slow_tool'),
+            ]
+        )
+
+    agent = Agent(FunctionModel(call_two_tools))
+
+    @agent.tool_plain
+    async def fast_failing_tool() -> str:
+        # Yield control so slow_tool can start, then raise.
+        await asyncio.sleep(0)
+        raise RuntimeError('boom')
+
+    @agent.tool_plain
+    async def slow_tool() -> str:
+        slow_tool_started.set()
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            slow_tool_cancelled.set()
+            raise
+        return 'done'  # pragma: no cover
+
+    tasks_before = asyncio.all_tasks()
+    with pytest.raises(RuntimeError, match='boom'):
+        await agent.run('call tools')
+
+    # Give the event loop a moment to process cancellations.
+    await asyncio.sleep(0)
+
+    # The slow tool must have started (confirming both tasks ran in parallel).
+    assert slow_tool_started.is_set(), 'slow_tool never started — not running in parallel'
+    # The slow tool must have been cancelled when fast_failing_tool raised.
+    assert slow_tool_cancelled.is_set(), 'slow_tool was not cancelled after RuntimeError'
+    # No new asyncio tasks should be left over from this run.
+    leaked = asyncio.all_tasks() - tasks_before
+    assert not leaked, f'Orphaned tasks remain: {leaked}'
+
+
 @pytest.mark.parametrize('mode', ['argument', 'contextmanager'])
 def test_sequential_calls(mode: Literal['argument', 'contextmanager']):
     """Test that tool calls are executed correctly when a `sequential` tool is present in the call."""