Skip to content

Commit bcac35a

Browse files
DeanChensjcopybara-github
authored andcommitted
fix(workflow): wait for all direct interrupts before rerunning rerun_on_resume nodes
When a `rerun_on_resume=True` node emits multiple interrupts simultaneously, `check_interception` previously triggered a rerun as soon as any single interrupt was resolved (`node.rerun_on_resume and recovered.resolved_ids`) because `_ChildScanState` did not distinguish interrupts raised directly by the node itself from descendant interrupts raised by dynamic child nodes. Track `direct_interrupt_ids` on `_ChildScanState` during event reconstruction and require `direct_unresolved` to be empty before rerunning a `rerun_on_resume` node when `unresolved` interrupts remain, while still allowing parent orchestrator nodes whose unresolved interrupts belong solely to descendant dynamic nodes to rerun and advance unblocked children. Co-authored-by: Shangjie Chen <[email protected]> PiperOrigin-RevId: 997167980
1 parent 531a86a commit bcac35a

4 files changed

Lines changed: 52 additions & 31 deletions

File tree

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,10 @@ class _ChildScanState:
7676
"""The agent the child asked to transfer to, or None if it asked for none."""
7777

7878
interrupt_ids: set[str] = field(default_factory=set)
79-
"""Every interrupt the child raised."""
79+
"""Every interrupt the child (or its descendants) raised."""
80+
81+
direct_interrupt_ids: set[str] = field(default_factory=set)
82+
"""Interrupts raised directly by the child node itself."""
8083

8184
resolved_ids: set[str] = field(default_factory=set)
8285
"""The subset of ``interrupt_ids`` a user response has come back for."""
@@ -425,6 +428,8 @@ def get_owner_key(event_path_builder: _NodePathBuilder) -> str | None:
425428
if interrupt_ids_to_process:
426429
for interrupt_id in interrupt_ids_to_process:
427430
child.interrupt_ids.add(interrupt_id)
431+
if is_direct:
432+
child.direct_interrupt_ids.add(interrupt_id)
428433
interrupt_owner[interrupt_id] = owner_key
429434

430435
schema_json = _extract_schema_from_event(event, interrupt_id)

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ def check_interception(
9090
)
9191

9292
unresolved = recovered.interrupt_ids - recovered.resolved_ids
93+
direct_unresolved = recovered.direct_interrupt_ids - recovered.resolved_ids
9394

9495
should_run = False
9596
output = None
@@ -99,9 +100,15 @@ def check_interception(
99100

100101
if unresolved:
101102
# Case 2: Cross-turn unresolved interrupts remain.
102-
# Rerun natively with resolved inputs if the node supports rerun and some
103-
# progress was made; otherwise remain waiting and bubble unresolved interrupts.
104-
if node.rerun_on_resume and recovered.resolved_ids:
103+
# A node with unresolved direct interrupts of its own must wait until all of
104+
# them are answered before rerunning. A parent node whose unresolved
105+
# interrupts belong only to descendant dynamic nodes reruns when progress was
106+
# made so the unblocked child can advance.
107+
if (
108+
not direct_unresolved
109+
and node.rerun_on_resume
110+
and recovered.resolved_ids
111+
):
105112
should_run = True
106113
resume_inputs = recovered.resolved_responses
107114
else:

‎tests/unittests/workflow/test_workflow_hitl.py‎

Lines changed: 4 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1053,23 +1053,7 @@ async def collect_events(agen):
10531053
assert node_b.received_inputs == ['final_output']
10541054

10551055

1056-
@pytest.mark.parametrize(
1057-
'resumable',
1058-
[
1059-
False,
1060-
pytest.param(
1061-
True,
1062-
marks=pytest.mark.xfail(
1063-
reason=(
1064-
'A rerun_on_resume node is re-run as soon as one interrupt'
1065-
' resolves instead of waiting for all pending interrupts;'
1066-
' fixing this needs a change to the shared replay'
1067-
' interception logic.'
1068-
)
1069-
),
1070-
),
1071-
],
1072-
)
1056+
@pytest.mark.parametrize('resumable', [False, True])
10731057
@pytest.mark.asyncio
10741058
async def test_rerun_on_resume_waits_for_all_interrupts(
10751059
request: pytest.FixtureRequest,
@@ -1182,14 +1166,10 @@ async def _run_impl(
11821166
)
11831167
assert final_state2['nodes']['NodeA']['interrupts'] == ['req2']
11841168

1185-
# The node should NOT have produced any RequestInput or data output in resumable mode.
1186-
# In non-resumable mode, it re-yields the pending interrupt 'req2'.
1169+
# The node should NOT have produced any RequestInput or data output since it
1170+
# did not rerun while req2 was still pending.
11871171
req_events2 = workflow_testing_utils.get_request_input_events(events2)
1188-
if resumable:
1189-
assert len(req_events2) == 0
1190-
else:
1191-
assert len(req_events2) == 1
1192-
assert get_request_input_interrupt_ids(req_events2[0]) == ['req2']
1172+
assert len(req_events2) == 0
11931173

11941174
# Run 3: provide req2 — now all interrupts resolved, node should rerun.
11951175
events3 = await runner.run_async(

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

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,12 +96,17 @@ def test_cross_turn_unresolved_interrupts_no_rerun():
9696
assert result.interrupts == {'fc-2'}
9797

9898

99-
def test_cross_turn_unresolved_interrupts_rerun():
100-
"""Cross-turn unresolved interrupts with rerun resolves progress and reruns."""
101-
# Given unresolved interrupts and node with rerun_on_resume
99+
def test_cross_turn_unresolved_descendant_interrupts_rerun():
100+
"""A rerun_on_resume node reruns once its own interrupts are all answered.
101+
102+
The remaining unresolved interrupt belongs to a descendant dynamic node, so
103+
the parent reruns to let that child advance.
104+
"""
105+
# Given the node's only direct interrupt resolved and a descendant's pending
102106
recovered = _ChildScanState(
103107
run_id='1',
104108
interrupt_ids={'fc-1', 'fc-2'},
109+
direct_interrupt_ids={'fc-1'},
105110
resolved_ids={'fc-1'},
106111
resolved_responses={'fc-1': 'ans'},
107112
)
@@ -118,6 +123,30 @@ def test_cross_turn_unresolved_interrupts_rerun():
118123
assert result.resume_inputs == {'fc-1': 'ans'}
119124

120125

126+
def test_cross_turn_unresolved_direct_interrupts_wait():
127+
"""A rerun_on_resume node waits while any of its own interrupts is pending."""
128+
# Given two direct interrupts of which only one has been answered
129+
recovered = _ChildScanState(
130+
run_id='1',
131+
interrupt_ids={'fc-1', 'fc-2'},
132+
direct_interrupt_ids={'fc-1', 'fc-2'},
133+
resolved_ids={'fc-1'},
134+
resolved_responses={'fc-1': 'ans'},
135+
)
136+
node = BaseNode(name='node', rerun_on_resume=True)
137+
138+
# When checked
139+
result = check_interception(
140+
node=node,
141+
recovered=recovered,
142+
)
143+
144+
# Then it keeps waiting on the unanswered one instead of rerunning
145+
assert not result.should_run
146+
assert result.interrupts == {'fc-2'}
147+
assert result.resume_inputs is None
148+
149+
121150
def test_cross_turn_completed():
122151
"""Cross-turn completed run fast-forwards output and route."""
123152
# Given a completed run from history

0 commit comments

Comments
 (0)