test(workflow): cover async cleanup before cancellation and retry

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,

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)