Skip to content

Commit a6ed859

Browse files
AkshayShah03copybara-github
authored andcommitted
fix: make JoinNode in a loop wait for each iteration's predecessors
Merge #7499 - A predecessor that routes away from the JoinNode no longer counts as arrived; the join waits until every incoming edge has fired for the current iteration. Previously this was order-dependent (the join fired if the routed-away predecessor happened to complete first), so some existing graphs with conditional edges into a JoinNode will see the join stop firing. - A node that completes without an output no longer leaves its previous iteration's output in node_outputs; downstream nodes (and the workflow's terminal output) now observe None for that run. - Known limitation: if an iteration restarts before the JoinNode fires (for example, when one predecessor arrives at the join while another routes back to restart the loop), the partial arrival carries over and a later firing can combine outputs from different iterations. Fixes #7498 Fixes #7515 PiperOrigin-RevId: 1001343757
1 parent 6a77d48 commit a6ed859

3 files changed

Lines changed: 235 additions & 2 deletions

File tree

‎docs/guides/workflow/join_node/index.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ validation_join = JoinNode(
112112

113113
- **The output is always a dictionary.** A `JoinNode` emits predecessor names mapped to their outputs and nothing else. When you need a different shape, transform it in the node downstream of the join.
114114
- **A predecessor on an untaken conditional path stops the join firing.** If a `JoinNode` declares a predecessor that sits on a conditional path, and that path is not taken, the join never triggers. The run does not hang waiting for it. The workflow finishes as soon as the branches that did run are done, and the join and everything after it is skipped silently, with no error anywhere. Every static predecessor a join declares in the graph has to execute before the join fires.
115+
- **In a loop, a `JoinNode` waits for every predecessor to arrive again in the current iteration.** If the loop restarts before the join fires, the arrivals already recorded carry over and the next firing can combine outputs from different iterations.
115116

116117
## Related samples
117118

‎src/google/adk/workflow/_workflow.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,14 @@ class _LoopState(DynamicNodeState):
9999
node_branches: dict[str, str] = field(default_factory=dict)
100100
"""Cached static node branches."""
101101

102+
join_arrivals: dict[str, set[str]] = field(default_factory=dict)
103+
"""Predecessors that arrived since each wait-for-all node last fired.
104+
105+
Keyed by the wait-for-all node's name. Cleared when that node fires, so in
106+
a loop it waits for every predecessor's next run instead of counting runs
107+
left over from an earlier iteration.
108+
"""
109+
102110
pending_tasks: dict[str, asyncio.Task[Context]] = field(default_factory=dict)
103111
"""Running static node tasks."""
104112

@@ -798,6 +806,8 @@ async def _handle_completion(
798806
node_state.resume_inputs.clear()
799807
if child_ctx.output is not None:
800808
loop_state.node_outputs[node_name] = child_ctx.output
809+
else:
810+
loop_state.node_outputs.pop(node_name, None)
801811
loop_state.node_branches[node_name] = (
802812
child_ctx._invocation_context.branch or ""
803813
)
@@ -824,22 +834,26 @@ def _buffer_barrier_trigger(
824834
) -> None:
825835
"""Buffer a trigger for target_name once all predecessors have completed.
826836
827-
No-op while any predecessor is still outstanding.
837+
No-op while any predecessor is still outstanding. Only runs completed
838+
since target_name last fired count, so each firing consumes them.
828839
"""
829840
graph = self._require_graph()
830841
predecessors = {
831842
e.from_node.name for e in graph.edges if e.to_node.name == target_name
832843
}
844+
arrivals = loop_state.join_arrivals.get(target_name, set())
833845
# START never executes, so it is satisfied as soon as the workflow begins.
834846
if not all(
835847
p == START.name
836848
or (
837-
loop_state.nodes.get(p)
849+
p in arrivals
850+
and loop_state.nodes.get(p)
838851
and loop_state.nodes[p].status == NodeStatus.COMPLETED
839852
)
840853
for p in predecessors
841854
):
842855
return
856+
loop_state.join_arrivals.pop(target_name, None)
843857

844858
outputs = {p: loop_state.node_outputs.get(p) for p in predecessors}
845859
branches = [loop_state.node_branches.get(p, "") for p in predecessors]
@@ -870,6 +884,7 @@ def _buffer_downstream_triggers(
870884
target_node = self._get_static_node_by_name(target_name)
871885

872886
if target_node._requires_all_predecessors:
887+
loop_state.join_arrivals.setdefault(target_name, set()).add(node_name)
873888
self._buffer_barrier_trigger(loop_state, target_name)
874889
else:
875890
# Normal node logic

‎tests/unittests/workflow/test_join_node.py‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,20 @@
1414

1515
"""Testings for the JoinNode."""
1616

17+
import asyncio
18+
1719
from google.adk import workflow
1820
from google.adk.apps import app
21+
from google.adk.apps.app import ResumabilityConfig
22+
from google.adk.events._interrupts import extract_event_interrupt_ids
23+
from google.adk.events.event import Event
24+
from google.adk.events.request_input import RequestInput
1925
from google.adk.workflow import _base_node as base_node
2026
from google.adk.workflow import _graph as workflow_graph
2127
from google.adk.workflow import _join_node as join_node
2228
from google.adk.workflow import START
2329
from google.adk.workflow._workflow import Workflow
30+
from google.adk.workflow.utils import _workflow_hitl_utils as hitl_utils
2431
from pydantic import BaseModel
2532
import pytest
2633

@@ -348,3 +355,213 @@ async def test_join_node_computes_common_branch_prefix(
348355
assert len(capture_events) > 0
349356
for e in capture_events:
350357
assert e.branch is None
358+
359+
360+
@pytest.mark.asyncio
361+
async def test_join_node_in_loop_waits_for_each_iteration(
362+
request: pytest.FixtureRequest,
363+
) -> None:
364+
"""Tests JoinNode in a loop waits for every predecessor's new run."""
365+
join_inputs: list[dict[str, str]] = []
366+
367+
def node_a() -> str:
368+
return 'A2' if join_inputs else 'A1'
369+
370+
def node_b() -> str:
371+
return 'B2' if join_inputs else 'B1'
372+
373+
def node_check(node_input: dict[str, str]) -> Event | None:
374+
join_inputs.append(node_input)
375+
return Event(route='again') if len(join_inputs) == 1 else None
376+
377+
join = join_node.JoinNode(name='join')
378+
agent = Workflow(
379+
name='wf',
380+
edges=[
381+
(START, node_a, node_b),
382+
((node_a, node_b), join, node_check),
383+
(node_check, {'again': node_a}),
384+
],
385+
)
386+
runner = testing_utils.InMemoryRunner(
387+
app=app.App(name=request.function.__name__, root_agent=agent)
388+
)
389+
390+
await runner.run_async('start')
391+
392+
assert join_inputs == [
393+
{'node_a': 'A1', 'node_b': 'B1'},
394+
{'node_a': 'A2', 'node_b': 'B2'},
395+
]
396+
397+
398+
@pytest.mark.asyncio
399+
@pytest.mark.parametrize('is_resumable', [False, True])
400+
async def test_join_node_in_loop_waits_for_predecessor_paused_for_input(
401+
request: pytest.FixtureRequest, is_resumable: bool
402+
) -> None:
403+
"""Tests JoinNode in a loop waits for a predecessor that asks the user."""
404+
join_inputs: list[dict[str, str]] = []
405+
406+
def node_a() -> str:
407+
return 'A2' if join_inputs else 'A1'
408+
409+
def node_b() -> RequestInput | str:
410+
return RequestInput(message='Continue?') if join_inputs else 'B1'
411+
412+
def node_check(node_input: dict[str, str]) -> Event | None:
413+
join_inputs.append(node_input)
414+
return Event(route='again') if len(join_inputs) == 1 else None
415+
416+
join = join_node.JoinNode(name='join')
417+
agent = Workflow(
418+
name='wf',
419+
edges=[
420+
(START, node_a, node_b),
421+
((node_a, node_b), join, node_check),
422+
(node_check, {'again': node_a}),
423+
],
424+
)
425+
runner = testing_utils.InMemoryRunner(
426+
app=app.App(
427+
name=request.function.__name__,
428+
root_agent=agent,
429+
resumability_config=ResumabilityConfig(is_resumable=is_resumable),
430+
)
431+
)
432+
433+
events = await runner.run_async('start')
434+
435+
assert join_inputs == [{'node_a': 'A1', 'node_b': 'B1'}]
436+
437+
request_events = workflow_testing_utils.get_request_input_events(events)
438+
[interrupt_id] = extract_event_interrupt_ids(request_events[0])
439+
await runner.run_async(
440+
new_message=hitl_utils.create_request_input_response(
441+
interrupt_id, {'result': 'B2'}
442+
),
443+
invocation_id=events[0].invocation_id,
444+
)
445+
446+
assert join_inputs == [
447+
{'node_a': 'A1', 'node_b': 'B1'},
448+
{'node_a': 'A2', 'node_b': 'B2'},
449+
]
450+
451+
452+
@pytest.mark.asyncio
453+
async def test_join_node_in_loop_clears_prior_output_when_node_returns_none(
454+
request: pytest.FixtureRequest,
455+
) -> None:
456+
"""Tests JoinNode receives None when a rerun predecessor returns None."""
457+
join_inputs: list[dict[str, str | None]] = []
458+
459+
def node_a() -> str | None:
460+
return None if join_inputs else 'A1'
461+
462+
def node_b() -> str:
463+
return 'B2' if join_inputs else 'B1'
464+
465+
def node_check(node_input: dict[str, str | None]) -> Event | None:
466+
join_inputs.append(node_input)
467+
return Event(route='again') if len(join_inputs) == 1 else None
468+
469+
join = join_node.JoinNode(name='join')
470+
agent = Workflow(
471+
name='wf',
472+
edges=[
473+
(START, node_a, node_b),
474+
((node_a, node_b), join, node_check),
475+
(node_check, {'again': node_a}),
476+
],
477+
)
478+
runner = testing_utils.InMemoryRunner(
479+
app=app.App(name=request.function.__name__, root_agent=agent)
480+
)
481+
482+
await runner.run_async('start')
483+
484+
assert join_inputs == [
485+
{'node_a': 'A1', 'node_b': 'B1'},
486+
{'node_a': None, 'node_b': 'B2'},
487+
]
488+
489+
490+
@pytest.mark.asyncio
491+
async def test_join_node_does_not_fire_when_predecessor_routes_away(
492+
request: pytest.FixtureRequest,
493+
) -> None:
494+
"""Tests JoinNode does not fire when a predecessor routes away from it."""
495+
496+
def node_a() -> Event:
497+
return Event(output='A1', route='away')
498+
499+
def node_b() -> str:
500+
return 'B1'
501+
502+
join = join_node.JoinNode(name='join')
503+
other_node = workflow_testing_utils.InputCapturingNode(name='other_node')
504+
capture = workflow_testing_utils.InputCapturingNode(name='capture')
505+
agent = Workflow(
506+
name='wf',
507+
edges=[
508+
(START, (node_a, node_b)),
509+
(node_a, {'to_join': join, 'away': other_node}),
510+
(node_b, join),
511+
(join, capture),
512+
],
513+
)
514+
runner = testing_utils.InMemoryRunner(
515+
app=app.App(name=request.function.__name__, root_agent=agent)
516+
)
517+
518+
await runner.run_async('start')
519+
520+
assert other_node.received_inputs == ['A1']
521+
assert capture.received_inputs == []
522+
523+
524+
@pytest.mark.asyncio
525+
async def test_join_node_waits_for_retriggered_in_flight_predecessor(
526+
request: pytest.FixtureRequest,
527+
) -> None:
528+
"""Tests JoinNode waits when an arrived predecessor has a run in flight."""
529+
a_runs = 0
530+
b_runs = 0
531+
532+
async def node_a() -> str:
533+
nonlocal a_runs
534+
a_runs += 1
535+
if a_runs == 2:
536+
for _ in range(5):
537+
await asyncio.sleep(0)
538+
return f'A{a_runs}'
539+
540+
def node_b() -> Event:
541+
nonlocal b_runs
542+
b_runs += 1
543+
route = 'retry' if b_runs == 1 else 'to_join'
544+
return Event(output=f'B{b_runs}', route=route)
545+
546+
def restart() -> None:
547+
return None
548+
549+
join = join_node.JoinNode(name='join')
550+
capture = workflow_testing_utils.InputCapturingNode(name='capture')
551+
agent = Workflow(
552+
name='wf',
553+
edges=[
554+
(START, (node_a, node_b)),
555+
(node_a, join),
556+
(node_b, {'to_join': join, 'retry': restart}),
557+
(restart, (node_a, node_b)),
558+
(join, capture),
559+
],
560+
)
561+
runner = testing_utils.InMemoryRunner(
562+
app=app.App(name=request.function.__name__, root_agent=agent)
563+
)
564+
565+
await runner.run_async('start')
566+
567+
assert capture.received_inputs == [{'node_a': 'A2', 'node_b': 'B2'}]

0 commit comments

Comments
 (0)