refactor(workflow): Make JoinNode stateless by moving join logic to orchestrator

boyangsvl · boyangsvl · commit cc3ba8a4cbce · 2026-04-16T14:10:25.000-07:00
Moved the stateful aggregation logic from JoinNode to the Workflow orchestrator to simplify node implementation and prevent state leaks in loops. Removed in_nodes from Context as it is no longer needed.

Change-Id: If0677f65c735ee3bce985478f7990b4a3335b7df
diff --git a/contributing/workflow_samples/fan_out_fan_in/tests/go.json b/contributing/workflow_samples/fan_out_fan_in/tests/go.json
@@ -57,56 +57,9 @@
       "output": "og"
     },
     {
-      "actions": {
-        "stateDelta": {
-          "root_agent@1/join_for_results@1_join_state": {
-            "make_uppercase": {
-              "branch": "make_uppercase@1",
-              "input": "GO"
-            }
-          }
-        }
-      },
       "author": "root_agent",
-      "branch": "make_uppercase@1",
       "id": "e-5",
       "invocationId": "i-1",
-      "nodeInfo": {
-        "path": "root_agent@1/join_for_results@1"
-      }
-    },
-    {
-      "actions": {
-        "stateDelta": {
-          "root_agent@1/join_for_results@1_join_state": {
-            "count_characters": {
-              "branch": "count_characters@1",
-              "input": 2
-            },
-            "make_uppercase": {
-              "branch": "make_uppercase@1",
-              "input": "GO"
-            }
-          }
-        }
-      },
-      "author": "root_agent",
-      "branch": "count_characters@1",
-      "id": "e-6",
-      "invocationId": "i-1",
-      "nodeInfo": {
-        "path": "root_agent@1/join_for_results@1"
-      }
-    },
-    {
-      "actions": {
-        "stateDelta": {
-          "root_agent@1/join_for_results@1_join_state": null
-        }
-      },
-      "author": "root_agent",
-      "id": "e-7",
-      "invocationId": "i-1",
       "nodeInfo": {
         "outputFor": [
           "root_agent@1/join_for_results@1"
@@ -129,7 +82,7 @@
         ],
         "role": "user"
       },
-      "id": "e-8",
+      "id": "e-6",
       "invocationId": "i-1",
       "nodeInfo": {
         "path": "root_agent@1/aggregate@1"
diff --git a/contributing/workflow_samples/nested_workflow/tests/1984.json b/contributing/workflow_samples/nested_workflow/tests/1984.json
@@ -36,7 +36,7 @@
       "content": {
         "parts": [
           {
-            "text": "LeBron James"
+            "text": "Scarlett Johansson"
           }
         ],
         "role": "model"
@@ -58,7 +58,7 @@
       "content": {
         "parts": [
           {
-            "text": "LeBron James, often hailed as \"King James,\" is a legendary basketball player widely considered one of the greatest of all time. With four NBA championships, four MVP awards, and the league's all-time scoring record, he has consistently shattered expectations and redefined player longevity. Beyond the court, James is a powerful voice for social justice and a successful entrepreneur and philanthropist through initiatives like the I PROMISE program."
+            "text": "Scarlett Johansson is an acclaimed actress renowned for her distinctive voice and versatile performances across a wide range of genres. From her captivating early roles in films like \"Lost in Translation\" to her iconic portrayal of Black Widow in the Marvel Cinematic Universe, she has consistently delivered powerful performances. A four-time Golden Globe nominee and two-time Academy Award nominee, she remains one of Hollywood's most bankable and respected stars."
           }
         ],
         "role": "model"
@@ -81,7 +81,7 @@
       "content": {
         "parts": [
           {
-            "text": "The Bhopal disaster occurred in December when a deadly leak of methyl isocyanate gas from a Union Carbide pesticide plant instantly killed thousands and sickened hundreds of thousands more. This catastrophic industrial accident remains one of the world's worst, profoundly influencing international environmental regulations and corporate accountability."
+            "text": "In December 1984, the Union Carbide chemical plant in Bhopal, India, experienced a catastrophic gas leak, releasing deadly methyl isocyanate. This disaster killed thousands instantly and caused hundreds of thousands of injuries, becoming one of the world's worst industrial accidents and a lasting symbol of corporate negligence."
           }
         ],
         "role": "model"
@@ -98,55 +98,31 @@
       }
     },
     {
-      "actions": {
-        "stateDelta": {
-          "root_agent@1/join_for_aggregation@1_join_state": {
-            "find_famous_person": {
-              "branch": "find_famous_person@1",
-              "input": "LeBron James, often hailed as \"King James,\" is a legendary basketball player widely considered one of the greatest of all time. With four NBA championships, four MVP awards, and the league's all-time scoring record, he has consistently shattered expectations and redefined player longevity. Beyond the court, James is a powerful voice for social justice and a successful entrepreneur and philanthropist through initiatives like the I PROMISE program."
-            }
-          }
-        }
-      },
       "author": "root_agent",
-      "branch": "find_famous_person@1",
       "id": "e-6",
       "invocationId": "i-1",
-      "nodeInfo": {
-        "path": "root_agent@1/join_for_aggregation@1"
-      }
-    },
-    {
-      "actions": {
-        "stateDelta": {
-          "root_agent@1/join_for_aggregation@1_join_state": null
-        }
-      },
-      "author": "root_agent",
-      "id": "e-7",
-      "invocationId": "i-1",
       "nodeInfo": {
         "outputFor": [
           "root_agent@1/join_for_aggregation@1"
         ],
         "path": "root_agent@1/join_for_aggregation@1"
       },
       "output": {
-        "find_famous_person": "LeBron James, often hailed as \"King James,\" is a legendary basketball player widely considered one of the greatest of all time. With four NBA championships, four MVP awards, and the league's all-time scoring record, he has consistently shattered expectations and redefined player longevity. Beyond the court, James is a powerful voice for social justice and a successful entrepreneur and philanthropist through initiatives like the I PROMISE program.",
-        "find_historical_event": "The Bhopal disaster occurred in December when a deadly leak of methyl isocyanate gas from a Union Carbide pesticide plant instantly killed thousands and sickened hundreds of thousands more. This catastrophic industrial accident remains one of the world's worst, profoundly influencing international environmental regulations and corporate accountability."
+        "find_famous_person": "Scarlett Johansson is an acclaimed actress renowned for her distinctive voice and versatile performances across a wide range of genres. From her captivating early roles in films like \"Lost in Translation\" to her iconic portrayal of Black Widow in the Marvel Cinematic Universe, she has consistently delivered powerful performances. A four-time Golden Globe nominee and two-time Academy Award nominee, she remains one of Hollywood's most bankable and respected stars.",
+        "find_historical_event": "In December 1984, the Union Carbide chemical plant in Bhopal, India, experienced a catastrophic gas leak, releasing deadly methyl isocyanate. This disaster killed thousands instantly and caused hundreds of thousands of injuries, becoming one of the world's worst industrial accidents and a lasting symbol of corporate negligence."
       }
     },
     {
       "author": "root_agent",
       "content": {
         "parts": [
           {
-            "text": "# Year: 1984\n\n## Famous Person Bio:\n\nLeBron James, often hailed as \"King James,\" is a legendary basketball player widely considered one of the greatest of all time. With four NBA championships, four MVP awards, and the league's all-time scoring record, he has consistently shattered expectations and redefined player longevity. Beyond the court, James is a powerful voice for social justice and a successful entrepreneur and philanthropist through initiatives like the I PROMISE program.\n\n## Historical Event:\n\nThe Bhopal disaster occurred in December when a deadly leak of methyl isocyanate gas from a Union Carbide pesticide plant instantly killed thousands and sickened hundreds of thousands more. This catastrophic industrial accident remains one of the world's worst, profoundly influencing international environmental regulations and corporate accountability."
+            "text": "# Year: 1984\n\n## Famous Person Bio:\n\nScarlett Johansson is an acclaimed actress renowned for her distinctive voice and versatile performances across a wide range of genres. From her captivating early roles in films like \"Lost in Translation\" to her iconic portrayal of Black Widow in the Marvel Cinematic Universe, she has consistently delivered powerful performances. A four-time Golden Globe nominee and two-time Academy Award nominee, she remains one of Hollywood's most bankable and respected stars.\n\n## Historical Event:\n\nIn December 1984, the Union Carbide chemical plant in Bhopal, India, experienced a catastrophic gas leak, releasing deadly methyl isocyanate. This disaster killed thousands instantly and caused hundreds of thousands of injuries, becoming one of the world's worst industrial accidents and a lasting symbol of corporate negligence."
           }
         ],
         "role": "user"
       },
-      "id": "e-8",
+      "id": "e-7",
       "invocationId": "i-1",
       "nodeInfo": {
         "path": "root_agent@1/aggregate_results@1"
diff --git a/src/google/adk/agents/context.py b/src/google/adk/agents/context.py
@@ -85,7 +85,6 @@ def __init__(
       node_path: str = '',
       run_id: str = '',
       triggered_by: str = '',
-      in_nodes: set[str] | None = None,
       resume_inputs: dict[str, Any] | None = None,
       schedule_dynamic_node_internal: ScheduleDynamicNodeInternal | None = None,
       node_rerun_on_resume: bool = True,
@@ -107,7 +106,6 @@ def __init__(
       node_path: The path of the current node in the workflow graph.
       run_id: The execution ID of the current node.
       triggered_by: The name of the node that triggered the current node.
-      in_nodes: Names of predecessor nodes.
       resume_inputs: Inputs for resuming node, keyed by interrupt id.
       node_rerun_on_resume: Whether the node reruns on resume.
       retry_count: Number of times this node has been retried.
@@ -134,9 +132,6 @@ def __init__(
     self._node_path = node_path
     self._run_id = run_id
     self._triggered_by = triggered_by
-    self._in_nodes = (
-        frozenset(in_nodes) if in_nodes is not None else frozenset()
-    )
     self._resume_inputs = resume_inputs or {}
     self._workflow_scheduler = schedule_dynamic_node_internal
     self._node_rerun_on_resume = node_rerun_on_resume
@@ -227,11 +222,6 @@ def attempt_count(self) -> int:
     """Returns the current attempt number (1-based)."""
     return self._attempt_count
 
-  @property
-  def in_nodes(self) -> frozenset[str]:
-    """Returns names of nodes that are predecessors of the current node."""
-    return self._in_nodes
-
   @property
   def resume_inputs(self) -> dict[str, Any]:
     """Returns inputs for resuming node, keyed by interrupt id."""
diff --git a/src/google/adk/workflow/_base_node.py b/src/google/adk/workflow/_base_node.py
@@ -243,6 +243,11 @@ async def _run_impl(
     )
     yield  # AsyncGenerator requires at least one yield statement
 
+  @property
+  def _requires_all_predecessors(self) -> bool:
+    """If True, the node waits for all predecessors to complete before running."""
+    return False
+
 
 START = BaseNode(name='__START__')
 """Sentinel node marking the entry point of a workflow graph.
diff --git a/src/google/adk/workflow/_join_node.py b/src/google/adk/workflow/_join_node.py
@@ -16,8 +16,8 @@
 
 from __future__ import annotations
 
-import logging
 from collections.abc import AsyncGenerator
+import logging
 from typing import Any
 
 from typing_extensions import override
@@ -48,10 +48,20 @@ class JoinNode(BaseNode):
   """A node that waits for all specified predecessors to trigger it before
   outputting."""
 
-  wait_for_output: bool = True
+  @property
+  @override
+  def _requires_all_predecessors(self) -> bool:
+    return True
 
-  def _get_state_key(self, node_path: str) -> str:
-    return f'{node_path}_join_state'
+  @override
+  def _validate_input_data(self, data: Any) -> Any:
+    """Validates individual trigger inputs against input_schema."""
+    if self.input_schema and isinstance(data, dict):
+      return {
+          k: self._validate_schema(v, self.input_schema)
+          for k, v in data.items()
+      }
+    return super()._validate_input_data(data)
 
   @override
   async def _run_impl(
@@ -60,59 +70,8 @@ async def _run_impl(
       ctx: Context,
       node_input: Any,
   ) -> AsyncGenerator[Any, None]:
-    if not ctx.in_nodes:
-      raise ValueError(
-          f'JoinNode {self.name} has no predecessors defined in graph.'
-      )
-
-    state_key = self._get_state_key(ctx.node_path)
-    join_state = (ctx.state.get(state_key) or {}).copy()
-
-    triggering_node = ctx.triggered_by
-    if not triggering_node:
-      logger.warning(
-          'JoinNode %s received trigger from node with no name. Ignoring.',
-          self.name,
-      )
-      return
-
-    if triggering_node not in ctx.in_nodes:
-      logger.warning(
-          'JoinNode %s received trigger from unexpected node %s. Ignoring.',
-          self.name,
-          triggering_node,
-      )
-      return
-
-    # Recording the output and branch from previous node.
-    join_state[triggering_node] = {
-        'input': node_input,
-        'branch': ctx._invocation_context.branch,
-    }
-
-    if set(join_state.keys()) == ctx.in_nodes:
-      # Extract outputs and branches
-      outputs = {}
-      branches = []
-      for k, v in join_state.items():
-        if isinstance(v, dict) and 'input' in v:
-          outputs[k] = v['input']
-          branches.append(v.get('branch') or '')
-        else:
-          # Fallback for old state structure
-          outputs[k] = v
-
-      common_branch = _get_common_branch_prefix(branches)
-
-      yield Event(
-          output=outputs,
-          branch=common_branch,
-          # Clear state for future runs
-          state={state_key: None},
-      )
-    else:
-      # Update state with recorded outputs from previous nodes and wait for
-      # more triggers.
-      yield Event(
-          state={state_key: join_state},
-      )
+    """JoinNode simply passes through the aggregated inputs provided by the orchestrator."""
+    yield Event(
+        output=node_input,
+        branch=ctx._invocation_context.branch,
+    )
diff --git a/src/google/adk/workflow/_node_runner.py b/src/google/adk/workflow/_node_runner.py
@@ -57,7 +57,6 @@ def __init__(
       run_id: str | None = None,
       # Graph context
       triggered_by: str = "",
-      in_nodes: set[str] | None = None,
       # Output delegation (use_as_output)
       additional_output_for_ancestor: str | None = None,
       # Resume state from a previous run
@@ -75,7 +74,6 @@ def __init__(
         counter string ("1", "2", …) unique per node path.
         Falls back to "1" if not provided.
       triggered_by: Name of the node that triggered this run.
-      in_nodes: Names of predecessor nodes in the graph.
       additional_output_for_ancestor: Ancestor node path whose
         output this node's output also represents (use_as_output).
       prior_output: Output from a previous run, carried
@@ -96,7 +94,6 @@ def __init__(
 
     # Graph context
     self._triggered_by = triggered_by
-    self._in_nodes = in_nodes
 
     # Output delegation
     self._additional_output_for_ancestor = additional_output_for_ancestor
@@ -246,7 +243,6 @@ def _create_child_context(
         schedule_dynamic_node_internal=scheduler,
         node_rerun_on_resume=self._node.rerun_on_resume,
         triggered_by=self._triggered_by,
-        in_nodes=self._in_nodes,
         output_for_ancestors=ancestors,
         event_author=self._parent_ctx.event_author,
         state_schema=self._node.state_schema
diff --git a/src/google/adk/workflow/_workflow.py b/src/google/adk/workflow/_workflow.py
diff --git a/tests/unittests/workflow/test_join_node.py b/tests/unittests/workflow/test_join_node.py
diff --git a/tests/unittests/workflow/test_workflow.py b/tests/unittests/workflow/test_workflow.py
