diff --git a/tests/unittests/workflow/test_workflow_failures.py b/tests/unittests/workflow/test_workflow_failures.py index 80df60adde5..fe3a78a79d8 100644 --- a/tests/unittests/workflow/test_workflow_failures.py +++ b/tests/unittests/workflow/test_workflow_failures.py @@ -775,6 +775,91 @@ def test_retry_config_normalizes_classes_to_strings(): assert config.exceptions == ['ValueError', 'KeyError'] +@pytest.mark.parametrize('topology', ['static', 'nested', 'dynamic']) +@pytest.mark.parametrize('termination', ['sibling_failure', 'caller_cancel']) +async def test_cancellation_awaits_cleanup_without_retrying( + topology: str, termination: str +) -> None: + """Cancellation finishes asynchronous child cleanup before returning. + + Setup: a retry-enabled worker holds a resource inside a static branch, + nested workflow, or dynamically scheduled child. + Act: fail a sibling or cancel the caller while the worker is running; + pause its asynchronous resource release until explicitly allowed. + Assert: the invocation stays pending until release completes, preserves + the original error/cancellation, and never retries the cancelled worker. + """ + started = asyncio.Event() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + attempts: list[int] = [] + released: list[str] = [] + + @node(retry_config=RetryConfig(max_attempts=3, initial_delay=0, jitter=0)) + async def worker(ctx: Context) -> str: + attempts.append(ctx.attempt_count) + started.set() + try: + await asyncio.Event().wait() + return 'unreachable' + finally: + cleanup_started.set() + await release_cleanup.wait() + released.append('resource') + + @node(rerun_on_resume=True) + async def parent(ctx: Context) -> Any: + return await ctx.run_node(worker) + + if topology == 'nested': + branch = Workflow(name='inner', edges=[(START, worker)]) + elif topology == 'dynamic': + branch = parent + else: + branch = worker + + @node + async def failing_sibling() -> str: + await started.wait() + raise ValueError('sibling failed') + + edges = [(START, branch)] + if termination == 'sibling_failure': + edges.append((START, failing_sibling)) + runner = testing_utils.InMemoryRunner(node=Workflow(name='wf', edges=edges)) + run_task = asyncio.create_task(runner.run_async('start')) + + try: + await asyncio.wait_for(started.wait(), timeout=5) + if termination == 'caller_cancel': + run_task.cancel() + await asyncio.wait_for(cleanup_started.wait(), timeout=5) + + # Keep cleanup blocked while giving an invocation that incorrectly skips + # the cleanup wait an opportunity to finish. Unlike wait_for, wait does + # not cancel the invocation when this observation window expires. + done, _ = await asyncio.wait({run_task}, timeout=0.1) + assert not done + assert released == [] + release_cleanup.set() + + if termination == 'caller_cancel': + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + else: + with pytest.raises(ValueError, match='sibling failed'): + await asyncio.wait_for(run_task, timeout=5) + + assert released == ['resource'] + assert attempts == [1] + finally: + # Unblock and reap the invocation even when a contract assertion fails. + release_cleanup.set() + if not run_task.done(): + run_task.cancel() + await asyncio.gather(run_task, return_exceptions=True) + + @pytest.mark.asyncio async def test_node_cancellation_on_sibling_failure( request: pytest.FixtureRequest, diff --git a/tests/unittests/workflow/test_workflow_node_timeout.py b/tests/unittests/workflow/test_workflow_node_timeout.py index b2a3368de3b..e83ae04c072 100644 --- a/tests/unittests/workflow/test_workflow_node_timeout.py +++ b/tests/unittests/workflow/test_workflow_node_timeout.py @@ -18,6 +18,7 @@ import asyncio +from google.adk.agents.context import Context from google.adk.apps import App from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService @@ -29,6 +30,7 @@ from google.genai import types import pytest +from .. import testing_utils from .._invariants import InvariantPlugin # --------------------------------------------------------------------------- @@ -202,3 +204,69 @@ async def slow_node(): await _run_workflow(outer_wf) assert 'inner_wf' in str(exc_info.value) + + +async def test_timeout_awaits_cleanup_before_retry() -> None: + """A timed-out attempt releases its resource before the next attempt starts. + + Setup: a worker with a timeout and one retry blocks on its first attempt. + Act: hold its asynchronous finalizer open after the deadline, then release it. + Assert: retry starts only after cleanup, receives attempt_count=2, and + produces the only downstream result. + """ + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + lifecycle: list[tuple[str, int]] = [] + received: list[str] = [] + + @node( + timeout=0.05, + retry_config=RetryConfig(max_attempts=2, initial_delay=0, jitter=0), + ) + async def worker(ctx: Context) -> str: + attempt = ctx.attempt_count + lifecycle.append(('start', attempt)) + try: + if attempt == 1: + await asyncio.Event().wait() + return 'recovered' + finally: + if attempt == 1: + cleanup_started.set() + await release_cleanup.wait() + lifecycle.append(('release', attempt)) + + @node + async def downstream(node_input: str) -> str: + received.append(node_input) + return f'processed: {node_input}' + + wf = Workflow(name='wf', edges=[(START, worker, downstream)]) + runner = testing_utils.InMemoryRunner(node=wf) + run_task = asyncio.create_task(runner.run_async('start')) + + try: + await asyncio.wait_for(cleanup_started.wait(), timeout=5) + # A timeout must not bypass the blocked finalizer to start the retry. + done, _ = await asyncio.wait({run_task}, timeout=0.1) + assert not done + assert lifecycle == [('start', 1)] + release_cleanup.set() + events = await asyncio.wait_for(run_task, timeout=5) + + assert lifecycle == [ + ('start', 1), + ('release', 1), + ('start', 2), + ('release', 2), + ] + assert received == ['recovered'] + assert [e.output for e in events if e.output is not None] == [ + 'recovered', + 'processed: recovered', + ] + finally: + release_cleanup.set() + if not run_task.done(): + run_task.cancel() + await asyncio.gather(run_task, return_exceptions=True)