Skip to content

Commit a4ba813

Browse files
DeanChensjcopybara-github
DeanChensj
authored and
copybara-github
committed
fix(workflow): buffer out-of-order completions in ReplaySequenceBarrier
During workflow resume, fast-forwarded nodes wait on `ReplaySequenceBarrier` while rerun nodes (such as `@node(rerun_on_resume=True)` pauses whose inputs were supplied, or nested `Workflow` nodes) execute live and call `check_and_advance` when they finish. Previously, `check_and_advance` ignored any completion key that arrived before the current head of `self.sequence`, causing downstream fast-forwarded nodes to time out after 15s with `WorkflowInvariantError: Replay divergence detected` whenever a rerun branch finished ahead of an earlier sibling in the recorded sequence. Track completed keys in `self._completed` and drain all contiguous completed keys at the head of `self.sequence` in `check_and_advance`. Closes #7502 Co-authored-by: Shangjie Chen <[email protected]> PiperOrigin-RevId: 999027711
1 parent 3800f06 commit a4ba813

3 files changed

Lines changed: 169 additions & 13 deletions

File tree

‎src/google/adk/workflow/utils/_replay_sequence_barrier.py‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ def __init__(self, sequence: list[str], timeout_sec: float = 15.0) -> None:
2828
self.sequence = sequence
2929
self.timeout_sec = timeout_sec
3030
self.current_index = 0
31+
self._completed: set[str] = set()
3132
self.events = {key: asyncio.Event() for key in sequence}
3233
if sequence:
3334
self.events[sequence[0]].set()
@@ -51,11 +52,13 @@ async def wait(self, key: str) -> None:
5152
) from e
5253

5354
def check_and_advance(self, key: str) -> None:
54-
"""Advance the sequence if the key matches the current expected execution."""
55-
if self.current_index < len(self.sequence):
56-
expected_key = self.sequence[self.current_index]
57-
if key == expected_key:
58-
self.current_index += 1
59-
if self.current_index < len(self.sequence):
60-
next_key = self.sequence[self.current_index]
61-
self.events[next_key].set()
55+
"""Advance the sequence through all completed keys at the head."""
56+
self._completed.add(key)
57+
while (
58+
self.current_index < len(self.sequence)
59+
and self.sequence[self.current_index] in self._completed
60+
):
61+
self.current_index += 1
62+
if self.current_index < len(self.sequence):
63+
next_key = self.sequence[self.current_index]
64+
self.events[next_key].set()

‎tests/unittests/workflow/test_workflow_hitl.py‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2316,3 +2316,148 @@ async def test_request_input_resume_after_earlier_invocation_completed(
23162316
isinstance(output, str) and output.startswith('resumed:')
23172317
for output in outputs3
23182318
), outputs3
2319+
2320+
2321+
@pytest.mark.parametrize('resumable', [False, True])
2322+
@pytest.mark.asyncio
2323+
async def test_parallel_request_input_answered_in_single_message_does_not_deadlock(
2324+
request: pytest.FixtureRequest,
2325+
resumable: bool,
2326+
):
2327+
"""Answering parallel RequestInput pauses in one message completes without barrier timeout."""
2328+
from google.adk.workflow import FunctionNode
2329+
from google.adk.workflow import JoinNode
2330+
2331+
def make_passthrough(name: str) -> FunctionNode:
2332+
def _run(node_input: Any) -> Any:
2333+
return node_input
2334+
2335+
return FunctionNode(func=_run, name=name)
2336+
2337+
def make_pause(name: str) -> FunctionNode:
2338+
async def _ask(ctx: Context, node_input: Any):
2339+
if name in ctx.resume_inputs:
2340+
yield Event(output=node_input)
2341+
return
2342+
yield RequestInput(interrupt_id=name, message=f'{name}?')
2343+
2344+
return FunctionNode(func=_ask, name=name, rerun_on_resume=True)
2345+
2346+
fork = make_passthrough('fork')
2347+
prepare = make_passthrough('prepare')
2348+
join = JoinNode(name='join')
2349+
approve = make_pause('approve')
2350+
confirm = make_pause('confirm')
2351+
step0 = make_passthrough('step0')
2352+
step1 = make_passthrough('step1')
2353+
step2 = make_passthrough('step2')
2354+
2355+
agent = Workflow(
2356+
name='parallel_hitl_wf',
2357+
edges=[
2358+
(START, fork),
2359+
(fork, [step0, prepare, confirm]),
2360+
(step0, step1),
2361+
(step1, step2),
2362+
(step2, join),
2363+
(prepare, approve),
2364+
(approve, join),
2365+
(confirm, join),
2366+
],
2367+
)
2368+
app = App(
2369+
name=request.function.__name__,
2370+
root_agent=agent,
2371+
resumability_config=(
2372+
ResumabilityConfig(is_resumable=True) if resumable else None
2373+
),
2374+
)
2375+
runner = testing_utils.InMemoryRunner(app=app)
2376+
2377+
events1 = await runner.run_async(testing_utils.get_user_content('{"x": 0}'))
2378+
req_events = workflow_testing_utils.get_request_input_events(events1)
2379+
assert len(req_events) == 2
2380+
invocation_id = events1[0].invocation_id
2381+
2382+
# Respond to both parallel interrupts in a single user message.
2383+
events2 = await runner.run_async(
2384+
new_message=types.UserContent(
2385+
parts=[
2386+
create_request_input_response('confirm', {'ok': True}),
2387+
create_request_input_response('approve', {'ok': True}),
2388+
]
2389+
),
2390+
invocation_id=invocation_id,
2391+
)
2392+
2393+
assert workflow_testing_utils.get_request_input_events(events2) == []
2394+
outputs2 = [e.output for e in events2 if e.output is not None]
2395+
assert outputs2
2396+
2397+
2398+
@pytest.mark.parametrize('resumable', [False, True])
2399+
@pytest.mark.asyncio
2400+
async def test_parallel_nested_workflows_with_hitl_resume_does_not_deadlock(
2401+
request: pytest.FixtureRequest,
2402+
resumable: bool,
2403+
):
2404+
"""Resuming parallel sub-workflows where one paused and the other has a downstream node completes."""
2405+
from google.adk.workflow import FunctionNode
2406+
from google.adk.workflow import JoinNode
2407+
2408+
def make_passthrough(name: str) -> FunctionNode:
2409+
def _run(node_input: Any) -> Any:
2410+
return node_input
2411+
2412+
return FunctionNode(func=_run, name=name)
2413+
2414+
async def _ask_gate(ctx: Context, node_input: Any):
2415+
if 'gate' in ctx.resume_inputs:
2416+
yield Event(output=node_input)
2417+
return
2418+
yield RequestInput(interrupt_id='gate', message='gate?')
2419+
2420+
lint = make_passthrough('lint')
2421+
test_node = make_passthrough('test')
2422+
checks = Workflow(name='checks', edges=[(START, lint), (lint, test_node)])
2423+
gate = FunctionNode(func=_ask_gate, name='gate', rerun_on_resume=True)
2424+
record = make_passthrough('record')
2425+
review = Workflow(name='review', edges=[(START, gate), (gate, record)])
2426+
fork = make_passthrough('fork')
2427+
report = make_passthrough('report')
2428+
join = JoinNode(name='join')
2429+
2430+
agent = Workflow(
2431+
name='parent_wf',
2432+
edges=[
2433+
(START, fork),
2434+
(fork, [checks, review]),
2435+
(checks, report),
2436+
(report, join),
2437+
(review, join),
2438+
],
2439+
)
2440+
app = App(
2441+
name=request.function.__name__,
2442+
root_agent=agent,
2443+
resumability_config=(
2444+
ResumabilityConfig(is_resumable=True) if resumable else None
2445+
),
2446+
)
2447+
runner = testing_utils.InMemoryRunner(app=app)
2448+
2449+
events1 = await runner.run_async(testing_utils.get_user_content('{"x": 0}'))
2450+
req_events = workflow_testing_utils.get_request_input_events(events1)
2451+
assert len(req_events) == 1
2452+
invocation_id = events1[0].invocation_id
2453+
2454+
events2 = await runner.run_async(
2455+
new_message=testing_utils.UserContent(
2456+
create_request_input_response('gate', {'ok': True})
2457+
),
2458+
invocation_id=invocation_id,
2459+
)
2460+
2461+
assert workflow_testing_utils.get_request_input_events(events2) == []
2462+
outputs2 = [e.output for e in events2 if e.output is not None]
2463+
assert outputs2

‎tests/unittests/workflow/utils/test_replay_sequence_barrier.py‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,17 +70,25 @@ async def wait_b():
7070
assert barrier.events['NodeB@1'].is_set()
7171

7272

73-
def test_barrier_advance_out_of_order_ignored():
74-
"""Verifies that out-of-order advance calls are ignored and do not progress index."""
75-
sequence = ['NodeA@1', 'NodeB@1']
73+
def test_barrier_advance_out_of_order_buffered_until_head_completes():
74+
"""Verifies out-of-order advance calls are buffered until earlier keys complete."""
75+
sequence = ['NodeA@1', 'NodeB@1', 'NodeC@1']
7676
barrier = ReplaySequenceBarrier(sequence)
7777

78-
# When second key tries to advance out of order
78+
# When second key advances out of order before the head key
7979
barrier.check_and_advance('NodeB@1')
8080

81-
# Then state remains unchanged
81+
# Then head index remains at 0 until NodeA@1 completes
8282
assert barrier.current_index == 0
8383
assert not barrier.events['NodeB@1'].is_set()
84+
assert not barrier.events['NodeC@1'].is_set()
85+
86+
# When the head key finally completes, the barrier drains through NodeB@1
87+
barrier.check_and_advance('NodeA@1')
88+
89+
assert barrier.current_index == 2
90+
assert barrier.events['NodeB@1'].is_set()
91+
assert barrier.events['NodeC@1'].is_set()
8492

8593

8694
@pytest.mark.asyncio

0 commit comments

Comments
 (0)