fix(workflow): make JoinNode in a loop wait for each iteration's predecessors

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]

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'},
+  ]