From 13781db2b4ea885bce77b4107610909f45ca24e5 Mon Sep 17 00:00:00 2001 From: cyhandsome-123~ Date: Mon, 19 Oct 2026 22:43:03 +0800 Subject: [PATCH] test(workflow): cover async cleanup before cancellation and retry Exercise sibling failures and caller cancellation across static, nested, and dynamic children. Verify that asynchronous cleanup finishes before the invocation returns and cancelled workers are not retried. Verify that a timed-out attempt releases its resources before the next attempt starts and only the recovered output reaches downstream nodes. Add Runner-based contract tests without modifying runtime behavior. --- .../workflow/test_workflow_failures.py | 85 +++++++++++++++++++ .../workflow/test_workflow_node_timeout.py | 68 +++++++++++++++ 2 files changed, 153 insertions(+) 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)