diff --git a/src/google/adk/flows/llm_flows/tools/_batch_executor.py b/src/google/adk/flows/llm_flows/tools/_batch_executor.py index a943fb19cd..5e5c62b414 100644 --- a/src/google/adk/flows/llm_flows/tools/_batch_executor.py +++ b/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() diff --git a/tests/unittests/flows/llm_flows/tools/test_batch_executor.py b/tests/unittests/flows/llm_flows/tools/test_batch_executor.py index 951c93db74..9e09dcda1a 100644 --- a/tests/unittests/flows/llm_flows/tools/test_batch_executor.py +++ b/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' ) diff --git a/tests/unittests/flows/llm_flows/tools/test_functions_parallel_error.py b/tests/unittests/flows/llm_flows/tools/test_functions_parallel_error.py index 3f1e2f2127..3be86bcdba 100644 --- a/tests/unittests/flows/llm_flows/tools/test_functions_parallel_error.py +++ b/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