chore: Rebase static graph dispatch onto dynamic scheduler

Use Context.run_node from within _workflow.py so static graphs also use the unified transfer loop.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 940728866

Author: DeanChensj

src/google/adk/workflow/_workflow.py

@@ -37,7 +37,6 @@
 from ._dynamic_node_scheduler import DynamicNodeState
 from ._graph import EdgeItem
 from ._graph import Graph
-from ._node_runner import NodeRunner
 from ._node_state import NodeState
 from ._node_status import NodeStatus
 from ._trigger import Trigger
@@ -536,7 +535,7 @@ def _start_node_task(
       node_name: str,
       trigger: Trigger,
   ) -> None:
-    """Create NodeRunner and start asyncio task for a node."""
+    """Start asyncio task for scheduling and executing a node."""
 
     assert self.graph is not None
 
@@ -602,22 +601,24 @@ async def return_ctx():
           key
       ].isolation_scope
 
-    runner = NodeRunner(
-        node=node,
-        parent_ctx=ctx,
-        run_id=run_id,
-        use_as_output=is_terminal,
-        use_sub_branch=trigger.use_sub_branch,
-        override_branch=trigger.branch,
-        override_isolation_scope=self._compute_isolation_scope_for_node(
-            node, trigger, ctx, run_id
-        ),
-    )
     resume_inputs = (
         dict(node_state.resume_inputs) if node_state.resume_inputs else None
     )
     loop_state.pending_tasks[node_name] = asyncio.create_task(
-        runner.run(node_input=trigger.input, resume_inputs=resume_inputs)
+        ctx._run_node_internal(
+            node,
+            node_input=trigger.input,
+            use_sub_branch=trigger.use_sub_branch,
+            override_branch=trigger.branch,
+            override_isolation_scope=self._compute_isolation_scope_for_node(
+                node, trigger, ctx, run_id
+            ),
+            return_ctx=True,
+            resume_inputs=resume_inputs,
+            run_id=run_id,
+            use_as_output=is_terminal,
+            skip_run_id_validation=True,
+        )
     )
 
   def _make_schedule_dynamic_node(

tests/unittests/workflow/test_workflow_hitl.py

@@ -1858,15 +1858,16 @@ async def _run_impl(
 async def test_multiple_invocations_isolation(request: pytest.FixtureRequest):
   """Verify that a new invocation ignores events from a previous invocation."""
 
+  run_counts = []
+
   class CounterNode(BaseNode):
     name: str = Field(default='counter_node')
-    run_count: int = Field(default=0)
 
     async def _run_impl(
         self, *, ctx: Context, node_input: Any
     ) -> AsyncGenerator[Any, None]:
-      self.run_count += 1
-      yield f'Run {self.run_count}'
+      run_counts.append(1)
+      yield f'Run {len(run_counts)}'
 
   node_a = CounterNode()
   wf = Workflow(name='wf', edges=[(START, node_a)])
@@ -1883,7 +1884,7 @@ async def _run_impl(
   ):
     events1.append(event)
 
-  assert node_a.run_count == 1
+  assert len(run_counts) == 1
 
   # Invocation 2 (New invocation in SAME session)
   msg2 = types.Content(parts=[types.Part(text='go 2')], role='user')
@@ -1894,7 +1895,7 @@ async def _run_impl(
     events2.append(event)
 
   # If isolation works, CounterNode should run AGAIN!
-  assert node_a.run_count == 2
+  assert len(run_counts) == 2
 
 
 @pytest.mark.asyncio

tests/unittests/workflow/test_workflow_live.py

@@ -27,6 +27,7 @@
 from google.adk.workflow._base_node import START
 from google.adk.workflow._workflow import Workflow
 from google.genai import types
+from pydantic import Field
 import pytest
 
 from . import testing_utils
@@ -39,6 +40,7 @@ class _MockNonLiveNode(BaseNode):
 
   called: bool = False
   actual_input: Any = None
+  shared_state: dict[str, Any] = Field(default_factory=dict)
 
   def __init__(self, *, name: str):
     super().__init__(name=name)
@@ -51,6 +53,8 @@ async def _run_impl(
   ) -> AsyncGenerator[Any, None]:
     self.called = True
     self.actual_input = node_input
+    self.shared_state["called"] = True
+    self.shared_state["actual_input"] = node_input
     yield Event(output=f"{self.name}_output")
 
 
@@ -72,29 +76,6 @@ async def _run_impl(
     yield Event(output=self.output_value)
 
 
-class _DynamicLiveSchedulerNode(BaseNode):
-  """A node that dynamically schedules a child live node using ctx.run_node()."""
-
-  child_node: BaseNode | None = None
-  child_output: Any = None
-
-  def __init__(self, *, name: str, child_node: BaseNode):
-    super().__init__(name=name, rerun_on_resume=True)
-    self.child_node = child_node
-
-  async def _run_impl(
-      self,
-      *,
-      ctx: Context,
-      node_input: Any,
-  ) -> AsyncGenerator[Any, None]:
-    if self.child_node:
-      self.child_output = await ctx.run_node(
-          self.child_node, node_input=node_input
-      )
-    yield Event(output=f"{self.name}_output")
-
-
 # --- Live Workflow Unit Tests (TDD) ---
 
 
@@ -468,6 +449,30 @@ async def test_nested_live_node_and_outer_live_node():
 @pytest.mark.asyncio
 async def test_dynamic_node_scheduling_of_live_node():
   """CUJ 4: A node in workflow dynamically schedules a live node using ctx.run_node()."""
+
+  class _DynamicLiveSchedulerNode(BaseNode):
+    """A node that dynamically schedules a child live node using ctx.run_node()."""
+
+    child_node: BaseNode | None = None
+    child_output: Any = None
+    shared_state: dict[str, Any] = Field(default_factory=dict)
+
+    def __init__(self, *, name: str, child_node: BaseNode):
+      super().__init__(name=name, rerun_on_resume=True)
+      self.child_node = child_node
+
+    async def _run_impl(
+        self,
+        *,
+        ctx: Context,
+        node_input: Any,
+    ) -> AsyncGenerator[Any, None]:
+      if self.child_node:
+        output = await ctx.run_node(self.child_node, node_input=node_input)
+        self.child_output = output
+        self.shared_state["child_output"] = output
+      yield Event(output=f"{self.name}_output")
+
   mock_model = testing_utils.MockModel.create(
       responses=[
           LlmResponse(
@@ -531,7 +536,9 @@ async def test_dynamic_node_scheduling_of_live_node():
       {"result": "DynamicLiveNode_output"},
       "SchedulerNode_output",
   ]
-  assert scheduler_node.child_output == {"result": "DynamicLiveNode_output"}
+  assert scheduler_node.shared_state.get("child_output") == {
+      "result": "DynamicLiveNode_output"
+  }
 
   # Assert content events
   content_texts = [
@@ -663,7 +670,7 @@ async def test_single_turn_agent_runs_as_non_live_in_live_session():
 
   outputs = [e.output for e in events if e.output is not None]
   assert outputs == ["initial_text_input", "capture_output"]
-  assert capture.actual_input == "SingleTurn_output"
+  assert capture.shared_state.get("actual_input") == "SingleTurn_output"
   # Verify that the model received the initial_text_input (node_input) and NOT the live queue audio
   assert len(mock_model.requests) == 1
   assert (

tests/unittests/workflow/test_workflow_schema.py

@@ -26,6 +26,7 @@
 from google.adk.workflow._workflow import Workflow
 from google.genai import types
 from pydantic import BaseModel
+from pydantic import ValidationError
 import pytest
 
 from .. import testing_utils
@@ -630,9 +631,10 @@ async def _run_impl(
   msg = types.Content(parts=[types.Part(text='hello')], role='user')
 
   # We expect it to raise ValidationError
-  from pydantic import ValidationError
-
-  with pytest.raises(ValidationError):
+  with pytest.raises(
+      ValueError,
+      match=r"Runtime schema validation failed for dynamic node 'node'",
+  ):
     async for event in runner.run_async(
         user_id='u', session_id=session.id, new_message=msg
     ):