fix(tools): cancel sibling tools on child-originated cancellation

src/google/adk/flows/llm_flows/tools/_batch_executor.py

@@ -160,7 +160,9 @@ async def _gather_or_cancel(tasks: list[asyncio.Task[_T]]) -> list[_T]:
   """Awaits every task, cancelling the rest as soon as one of them fails."""
   try:
     return list(await asyncio.gather(*tasks))
-  except Exception:
+  except (Exception, asyncio.CancelledError):
+    # CancelledError is BaseException, so Exception alone would leave
+    # siblings running after a child-originated cancellation.
     for t in tasks:
       if not t.done():
         t.cancel()

tests/unittests/flows/llm_flows/tools/test_batch_executor.py

@@ -146,6 +146,77 @@ async def slow_worker() -> None:
   assert cancelled
 
 
+@pytest.mark.asyncio
+async def test_gather_or_cancel_cancels_siblings_on_cancelled_error() -> None:
+  """A child-raised CancelledError cancels unfinished sibling tasks."""
+  started = asyncio.Event()
+  sibling_cancelled = False
+  sibling_finished = False
+
+  async def failing_worker() -> None:
+    await started.wait()
+    raise asyncio.CancelledError()
+
+  async def slow_worker() -> None:
+    nonlocal sibling_cancelled, sibling_finished
+    started.set()
+    try:
+      await asyncio.sleep(10)
+      sibling_finished = True
+    except asyncio.CancelledError:
+      sibling_cancelled = True
+      raise
+
+  tasks = [
+      asyncio.create_task(failing_worker()),
+      asyncio.create_task(slow_worker()),
+  ]
+  with pytest.raises(asyncio.CancelledError):
+    await _batch_tool_executor._gather_or_cancel(tasks)
+
+  assert sibling_cancelled
+  assert not sibling_finished
+  assert all(t.done() for t in tasks)
+
+
+@pytest.mark.asyncio
+async def test_gather_or_cancel_cancels_siblings_when_child_cancels_itself() -> (
+    None
+):
+  """A child that cancels its own task still tears down unfinished siblings."""
+  started = asyncio.Event()
+  sibling_cancelled = False
+  sibling_finished = False
+
+  async def self_cancelling_worker() -> None:
+    await started.wait()
+    task = asyncio.current_task()
+    assert task is not None
+    task.cancel()
+    await asyncio.sleep(0)
+
+  async def slow_worker() -> None:
+    nonlocal sibling_cancelled, sibling_finished
+    started.set()
+    try:
+      await asyncio.sleep(10)
+      sibling_finished = True
+    except asyncio.CancelledError:
+      sibling_cancelled = True
+      raise
+
+  tasks = [
+      asyncio.create_task(self_cancelling_worker()),
+      asyncio.create_task(slow_worker()),
+  ]
+  with pytest.raises(asyncio.CancelledError):
+    await _batch_tool_executor._gather_or_cancel(tasks)
+
+  assert sibling_cancelled
+  assert not sibling_finished
+  assert all(t.done() for t in tasks)
+
+
 _probe: contextvars.ContextVar[str] = contextvars.ContextVar(
     'probe', default='unset'
 )

tests/unittests/flows/llm_flows/tools/test_functions_parallel_error.py

@@ -13,10 +13,14 @@
 # limitations under the License.
 
 import asyncio
+from contextlib import aclosing
 from typing import Any
 
 from google.adk.agents.llm_agent import Agent
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.run_config import StreamingMode
 from google.adk.flows.llm_flows import functions
+from google.adk.live import LiveRequestQueue
 from google.adk.tools.tool_context import ToolContext
 from google.genai import types
 import pytest
@@ -84,3 +88,163 @@ async def sleep_tool(tool_context: ToolContext) -> str:
   assert sleep_started
   assert not sleep_completed
   assert sleep_cancelled
+
+
+def _parallel_tool_responses() -> list[list[types.Part]]:
+  return [
+      [
+          function_call('id_1', 'stopper', {}),
+          function_call('id_2', 'sibling', {}),
+      ],
+      [
+          types.Part.from_text(text='final response'),
+      ],
+  ]
+
+
+async def _run_parallel_child_cancel(
+    stopper,
+    sibling,
+    execution: str = 'run_async',
+) -> None:
+  agent = Agent(
+      name='root_agent',
+      model=testing_utils.MockModel.create(
+          responses=_parallel_tool_responses()
+      ),
+      tools=[stopper, sibling],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  session = runner.session
+  user_message = types.Content(
+      role='user', parts=[types.Part(text='test')]
+  )
+  # Runner swallows CancelledError at root-task cleanup so the caller is
+  # not itself cancelled; the invariant is that sibling tools are torn down
+  # before that iterator returns.
+  if execution == 'live':
+    live_queue = LiveRequestQueue()
+    live_queue.send_content(user_message)
+    live_queue.close()
+
+    async def _consume_live() -> None:
+      async with aclosing(
+          runner.runner.run_live(
+              user_id=session.user_id,
+              session_id=session.id,
+              live_request_queue=live_queue,
+              run_config=RunConfig(response_modalities=['TEXT']),
+          )
+      ) as agen:
+        async for _ in agen:
+          pass
+
+    await asyncio.wait_for(_consume_live(), timeout=10)
+    return
+
+  streaming_mode = (
+      StreamingMode.SSE if execution == 'sse' else StreamingMode.NONE
+  )
+  async with aclosing(
+      runner.runner.run_async(
+          user_id=session.user_id,
+          session_id=session.id,
+          new_message=user_message,
+          run_config=RunConfig(streaming_mode=streaming_mode),
+      )
+  ) as agen:
+    async for _ in agen:
+      pass
+
+
+_EXECUTIONS = [
+    pytest.param('run_async', id='run-async'),
+    pytest.param('sse', id='run-async-sse'),
+    pytest.param('live', id='run-live'),
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('execution', _EXECUTIONS)
+async def test_parallel_function_call_cancels_siblings_on_cancelled_error(
+    execution: str,
+):
+  """A tool that raises CancelledError cancels unfinished sibling tools."""
+  started = asyncio.Event()
+  release = asyncio.Event()
+  sibling_task = None
+  sibling_completed = False
+  sibling_cancelled = False
+
+  async def stopper(tool_context: ToolContext) -> str:
+    await started.wait()
+    raise asyncio.CancelledError()
+
+  async def sibling(tool_context: ToolContext) -> str:
+    nonlocal sibling_task, sibling_completed, sibling_cancelled
+    sibling_task = asyncio.current_task()
+    started.set()
+    try:
+      await release.wait()
+      sibling_completed = True
+      return 'ok'
+    except asyncio.CancelledError:
+      sibling_cancelled = True
+      raise
+
+  await _run_parallel_child_cancel(stopper, sibling, execution=execution)
+
+  pending = sibling_task is not None and not sibling_task.done()
+  release.set()
+  if sibling_task is not None:
+    await asyncio.gather(sibling_task, return_exceptions=True)
+
+  assert sibling_task is not None
+  assert sibling_cancelled
+  assert not sibling_completed
+  assert not pending
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('execution', _EXECUTIONS)
+async def test_parallel_function_call_cancels_siblings_when_tool_cancels_itself(
+    execution: str,
+):
+  """A tool that cancels its own task cancels unfinished sibling tools."""
+  started = asyncio.Event()
+  release = asyncio.Event()
+  sibling_task = None
+  sibling_completed = False
+  sibling_cancelled = False
+
+  async def stopper(tool_context: ToolContext) -> str:
+    await started.wait()
+    task = asyncio.current_task()
+    assert task is not None
+    task.cancel()
+    await asyncio.sleep(0)
+    return 'should not reach'
+
+  async def sibling(tool_context: ToolContext) -> str:
+    nonlocal sibling_task, sibling_completed, sibling_cancelled
+    sibling_task = asyncio.current_task()
+    started.set()
+    try:
+      await release.wait()
+      sibling_completed = True
+      return 'ok'
+    except asyncio.CancelledError:
+      sibling_cancelled = True
+      raise
+
+  await _run_parallel_child_cancel(stopper, sibling, execution=execution)
+
+  pending = sibling_task is not None and not sibling_task.done()
+  release.set()
+  if sibling_task is not None:
+    await asyncio.gather(sibling_task, return_exceptions=True)
+
+  assert sibling_task is not None
+  assert sibling_cancelled
+  assert not sibling_completed
+  assert not pending