From 1fe1a1a8ab45871e6df9d0ecfb9ea0ea8881a819 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Sat, 10 Oct 2026 18:03:59 -0500 Subject: [PATCH] fix(workflow): make JoinNode in a loop wait for each iteration's predecessors _buffer_barrier_trigger treated a wait-for-all node as ready when every predecessor's status was COMPLETED. That status is never consumed when the join fires, so in a loop the join fired as soon as the first predecessor of the next iteration completed, using the other predecessors' outputs from the previous iteration, and then fired again when they arrived. Track the predecessors that completed since each wait-for-all node last fired and clear them when it fires. Resume replays recorded completions through the same path, so the tracking is rebuilt without persisting anything new. Fixes #7498 --- src/google/adk/workflow/_workflow.py | 37 +++-- tests/unittests/workflow/test_join_node.py | 158 +++++++++++++++++++++ 2 files changed, 186 insertions(+), 9 deletions(-) diff --git a/src/google/adk/workflow/_workflow.py b/src/google/adk/workflow/_workflow.py index aea1ace156f..bc2e3c2c960 100644 --- a/src/google/adk/workflow/_workflow.py +++ b/src/google/adk/workflow/_workflow.py @@ -97,6 +97,14 @@ class _LoopState(DynamicNodeState): node_branches: dict[str, str] = field(default_factory=dict) """Cached static node branches.""" + join_arrivals: dict[str, set[str]] = field(default_factory=dict) + """Predecessors that completed since each wait-for-all node last fired. + + Keyed by the wait-for-all node's name. Cleared when that node fires, so in + a loop it waits for every predecessor's next run instead of counting runs + left over from an earlier iteration. + """ + pending_tasks: dict[str, asyncio.Task[Context]] = field(default_factory=dict) """Running static node tasks.""" @@ -790,6 +798,7 @@ async def _handle_completion( loop_state.node_branches[node_name] = ( child_ctx._invocation_context.branch or "" ) + self._record_join_arrivals(loop_state, node_name) # A genuine execution records a checkpoint on completion. A replayed # fast-forward instead re-surfaces its recovered output so a resumable @@ -808,27 +817,37 @@ async def _handle_completion( child_ctx._invocation_context.branch, ) + def _record_join_arrivals( + self, loop_state: _LoopState, node_name: str + ) -> None: + """Record node_name's completed run for each wait-for-all successor.""" + graph = self._require_graph() + for edge in graph.edges: + if ( + edge.from_node.name == node_name + and edge.to_node._requires_all_predecessors + ): + loop_state.join_arrivals.setdefault(edge.to_node.name, set()).add( + node_name + ) + def _buffer_barrier_trigger( self, loop_state: _LoopState, target_name: str ) -> None: """Buffer a trigger for target_name once all predecessors have completed. - No-op while any predecessor is still outstanding. + No-op while any predecessor is still outstanding. Only runs completed + since target_name last fired count, so each firing consumes them. """ graph = self._require_graph() predecessors = { e.from_node.name for e in graph.edges if e.to_node.name == target_name } + arrivals = loop_state.join_arrivals.get(target_name, set()) # START never executes, so it is satisfied as soon as the workflow begins. - if not all( - p == START.name - or ( - loop_state.nodes.get(p) - and loop_state.nodes[p].status == NodeStatus.COMPLETED - ) - for p in predecessors - ): + if not all(p == START.name or p in arrivals for p in predecessors): return + loop_state.join_arrivals.pop(target_name, None) outputs = {p: loop_state.node_outputs.get(p) for p in predecessors} branches = [loop_state.node_branches.get(p, "") for p in predecessors] diff --git a/tests/unittests/workflow/test_join_node.py b/tests/unittests/workflow/test_join_node.py index 721454a56d8..0d7af700788 100644 --- a/tests/unittests/workflow/test_join_node.py +++ b/tests/unittests/workflow/test_join_node.py @@ -16,11 +16,16 @@ from google.adk import workflow from google.adk.apps import app +from google.adk.apps.app import ResumabilityConfig +from google.adk.events.event import Event +from google.adk.events.request_input import RequestInput from google.adk.workflow import _base_node as base_node from google.adk.workflow import _graph as workflow_graph from google.adk.workflow import _join_node as join_node from google.adk.workflow import START from google.adk.workflow._workflow import Workflow +from google.adk.workflow.utils._workflow_hitl_utils import create_request_input_response +from google.adk.workflow.utils._workflow_hitl_utils import get_request_input_interrupt_ids from pydantic import BaseModel import pytest @@ -348,3 +353,156 @@ async def test_join_node_computes_common_branch_prefix( assert len(capture_events) > 0 for e in capture_events: assert e.branch is None + + +@pytest.mark.asyncio +async def test_join_node_in_loop_waits_for_each_iteration( + request: pytest.FixtureRequest, +): + """Tests JoinNode in a loop waits for every predecessor's new run. + + NodeB sits one hop further from START than NodeA, so on the second + iteration NodeA completes first. The join must not fire with NodeB's + output left over from the first iteration. + """ + iteration = 0 + join_inputs = [] + + def node_a(node_input): + nonlocal iteration + iteration += 1 + return f'A{iteration}' + + def node_x(node_input): + return 'x' + + def node_b(node_input): + return f'B{iteration}' + + def node_check(node_input): + join_inputs.append(node_input) + if len(join_inputs) < 2: + yield Event(route='again') + else: + yield Event(route='done', output='finished') + + def node_end(node_input): + return node_input + + a = workflow.node(node_a, name='NodeA') + x = workflow.node(node_x, name='NodeX') + b = workflow.node(node_b, name='NodeB') + check = workflow.node(node_check, name='NodeCheck') + end = workflow.node(node_end, name='NodeEnd') + join = join_node.JoinNode(name='NodeJoin') + agent = workflow.Workflow( + name='test_join_node_loop', + edges=[ + workflow_graph.Edge(from_node=START, to_node=a), + workflow_graph.Edge(from_node=START, to_node=x), + workflow_graph.Edge(from_node=x, to_node=b), + workflow_graph.Edge(from_node=a, to_node=join), + workflow_graph.Edge(from_node=b, to_node=join), + workflow_graph.Edge(from_node=join, to_node=check), + workflow_graph.Edge(from_node=check, to_node=a, route='again'), + workflow_graph.Edge(from_node=check, to_node=x, route='again'), + workflow_graph.Edge(from_node=check, to_node=end, route='done'), + ], + ) + runner = testing_utils.InMemoryRunner( + app=app.App(name=request.function.__name__, root_agent=agent) + ) + + await runner.run_async(testing_utils.get_user_content('start')) + + assert join_inputs == [ + {'NodeA': 'A1', 'NodeB': 'B1'}, + {'NodeA': 'A2', 'NodeB': 'B2'}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize('is_resumable', [False, True]) +async def test_join_node_in_loop_waits_for_predecessor_paused_for_input( + request: pytest.FixtureRequest, is_resumable: bool +): + """Tests JoinNode in a loop waits for a predecessor that asks the user. + + On the second iteration NodeB pauses for input. The join must not fire + with NodeB's first-iteration output while it waits, and must fire once + with both second-iteration outputs after the resume. + """ + a_runs = 0 + b_runs = 0 + join_inputs = [] + + def node_a(node_input): + nonlocal a_runs + a_runs += 1 + return f'A{a_runs}' + + def node_x(node_input): + return 'x' + + def node_b(ctx, node_input): + nonlocal b_runs + if b_runs == 1 and 'ask' not in ctx.resume_inputs: + return RequestInput(interrupt_id='ask', message='Continue?') + b_runs += 1 + return f'B{b_runs}' + + def node_check(node_input): + join_inputs.append(node_input) + if len(join_inputs) < 2: + yield Event(route='again') + else: + yield Event(route='done', output='finished') + + def node_end(node_input): + return node_input + + a = workflow.node(node_a, name='NodeA') + x = workflow.node(node_x, name='NodeX') + b = workflow.node(node_b, name='NodeB', rerun_on_resume=True) + check = workflow.node(node_check, name='NodeCheck') + end = workflow.node(node_end, name='NodeEnd') + join = join_node.JoinNode(name='NodeJoin') + agent = workflow.Workflow( + name='test_join_node_loop_hitl', + edges=[ + workflow_graph.Edge(from_node=START, to_node=a), + workflow_graph.Edge(from_node=START, to_node=x), + workflow_graph.Edge(from_node=x, to_node=b), + workflow_graph.Edge(from_node=a, to_node=join), + workflow_graph.Edge(from_node=b, to_node=join), + workflow_graph.Edge(from_node=join, to_node=check), + workflow_graph.Edge(from_node=check, to_node=a, route='again'), + workflow_graph.Edge(from_node=check, to_node=x, route='again'), + workflow_graph.Edge(from_node=check, to_node=end, route='done'), + ], + ) + runner = testing_utils.InMemoryRunner( + app=app.App( + name=request.function.__name__, + root_agent=agent, + resumability_config=ResumabilityConfig(is_resumable=is_resumable), + ) + ) + + events = await runner.run_async(testing_utils.get_user_content('start')) + + assert join_inputs == [{'NodeA': 'A1', 'NodeB': 'B1'}] + + request_events = workflow_testing_utils.get_request_input_events(events) + interrupt_id = get_request_input_interrupt_ids(request_events[0])[0] + await runner.run_async( + new_message=testing_utils.UserContent( + create_request_input_response(interrupt_id, {'ok': True}) + ), + invocation_id=events[0].invocation_id, + ) + + assert join_inputs == [ + {'NodeA': 'A1', 'NodeB': 'B1'}, + {'NodeA': 'A2', 'NodeB': 'B2'}, + ]