refactor: Switch runners._drive_root_node to use context.run_node

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

Author: DeanChensj

src/google/adk/agents/context.py

@@ -491,6 +491,7 @@ async def _run_node_internal(
       raise_on_wait: bool = False,
       return_ctx: bool = False,
       resume_inputs: dict[str, Any] | None = None,
+      skip_run_id_validation: bool = False,
   ) -> Any:
     """Executes a node dynamically (Internal Orchestration API).
 
@@ -516,102 +517,165 @@ async def _run_node_internal(
     if isinstance(node, BaseAgent) and isinstance(built_node, BaseAgent):
       built_node.parent_agent = node.parent_agent
 
-    # Mode 1: Running within a Workflow graph.
-    # The workflow orchestrator provides a scheduler to handle resume, dedup,
-    # etc.
-    if self._workflow_scheduler:
-      from ..workflow._errors import NodeInterruptedError
+    # Output delegation: once set, the calling node's own output
+    # events are suppressed — the child's output (annotated with
+    # output_for) becomes the calling node's output.
+    # We validate and set this upfront before entering the loop.
+    if use_as_output:
+      from ..workflow._workflow import Workflow
 
-      # Output delegation: once set, the calling node's own output
-      # events are suppressed — the child's output (annotated with
-      # output_for) becomes the calling node's output.
-      if use_as_output:
+      if not isinstance(self.node, Workflow):
         if self._output_delegated:
           raise ValueError(
               f'Node {self.node_path} already has a use_as_output delegate.'
           )
         self._output_delegated = True
 
-      if run_id:
-        if run_id.isdigit():
-          raise ValueError(
-              f'Explicit run_id "{run_id}" for node "{built_node.name}" must'
-              ' contain non-numeric characters to prevent collision with'
-              ' auto-generated IDs.'
+    # Pointers to track the active execution state in the transfer loop.
+    # These will be updated dynamically if an agent transfers execution.
+    curr_parent_ctx = self
+    curr_node = built_node
+    curr_run_id = run_id
+    curr_input = node_input
+
+    # Active Execution Loop: Handles both standard execution and sequential Agent Transfers
+    # (e.g. Agent A transferring to Agent B). Instead of recursive execution, we use this
+    # loop to execute the target agent in-place, updating pointers and 'continuing' the loop.
+    while True:
+      curr_use_as_output = use_as_output if (curr_parent_ctx is self) else False
+      if self._workflow_scheduler:
+        # --- Mode 1: Workflow Execution ---
+        # The node is running as part of a Workflow graph. We must delegate execution
+        # to the workflow scheduler to handle graph dependencies and state.
+        from ..workflow._errors import NodeInterruptedError
+
+        # Validate or auto-generate run_id for this scheduler execution.
+        if curr_run_id:
+          if curr_run_id.isdigit() and not skip_run_id_validation:
+            raise ValueError(
+                f'Explicit run_id "{curr_run_id}" for node "{curr_node.name}"'
+                ' must contain non-numeric characters to prevent collision'
+                ' with auto-generated IDs.'
+            )
+        elif not curr_run_id:
+          curr_parent_ctx._child_run_counters[curr_node.name] = (
+              curr_parent_ctx._child_run_counters.get(curr_node.name, 0) + 1
           )
+          curr_run_id = str(curr_parent_ctx._child_run_counters[curr_node.name])
+
+        child_ctx = await curr_parent_ctx._workflow_scheduler(
+            curr_parent_ctx,
+            curr_node,
+            curr_input,
+            node_name=curr_node.name,
+            use_as_output=curr_use_as_output,
+            run_id=curr_run_id,
+            use_sub_branch=use_sub_branch,
+            override_branch=override_branch,
+            override_isolation_scope=override_isolation_scope,
+        )
       else:
-        self._child_run_counters[built_node.name] = (
-            self._child_run_counters.get(built_node.name, 0) + 1
+        # --- Mode 2: Standalone Execution ---
+        # The node is running independently (outside of a workflow).
+        # We run it directly using NodeRunner.
+        child_ctx = await curr_parent_ctx._run_node_standalone(
+            curr_node,
+            curr_input,
+            use_as_output=curr_use_as_output,
+            use_sub_branch=use_sub_branch,
+            override_branch=override_branch,
+            override_isolation_scope=override_isolation_scope,
+            run_id=curr_run_id,
+            resume_inputs=resume_inputs,
         )
-        run_id = str(self._child_run_counters[built_node.name])
-
-      child_ctx = await self._workflow_scheduler(
-          self,
-          built_node,
-          node_input,
-          node_name=built_node.name,
-          use_as_output=use_as_output,
-          run_id=run_id,
-          use_sub_branch=use_sub_branch,
-          override_branch=override_branch,
-          override_isolation_scope=override_isolation_scope,
+
+      # Extract the transfer target if the node requested an agent transfer.
+      transfer_to_agent = (
+          child_ctx.actions.transfer_to_agent if child_ctx else None
       )
-      if child_ctx.error:
-        from ..workflow._errors import DynamicNodeFailError
 
-        raise DynamicNodeFailError(
-            message=f'Dynamic node {built_node.name} failed',
-            error=child_ctx.error,
-            error_node_path=child_ctx.error_node_path,
+      # Post-Execution Validation: If the caller expects the raw output (not the Context),
+      # we check for errors or interrupts and raise them immediately.
+      if not return_ctx:
+        if child_ctx.error:
+          from ..workflow._errors import DynamicNodeFailError
+
+          raise DynamicNodeFailError(
+              message=f'Dynamic node {curr_node.name} failed',
+              error=child_ctx.error,
+              error_node_path=child_ctx.error_node_path,
+          )
+        if child_ctx.interrupt_ids:
+          from ..workflow._errors import NodeInterruptedError
+
+          # Propagate child's interrupt_ids to this node's ctx
+          # so NodeRunner sees them after catching the error.
+          curr_parent_ctx._interrupt_ids.update(child_ctx.interrupt_ids)
+          raise NodeInterruptedError()
+        # When the caller passes raise_on_wait=True, surface a child
+        # that's WAITING (wait_for_output, no output, not transferring)
+        # as NodeInterruptedError so the parent's NodeRunner records
+        # the parent as WAITING instead of falsely COMPLETED.
+        if (
+            raise_on_wait
+            and curr_node.wait_for_output
+            and child_ctx.output is None
+            and not transfer_to_agent
+        ):
+          from ..workflow._errors import NodeInterruptedError
+
+          raise NodeInterruptedError()
+
+      # Handle Agent Transfer: If a transfer was requested, we resolve the target agent
+      # and its parent context, update loop pointers, and continue to the next iteration.
+      if isinstance(transfer_to_agent, str):
+        target_name = transfer_to_agent
+        root_agent = getattr(curr_node, 'root_agent', None)
+        if not root_agent:
+          raise ValueError(f'Cannot find root_agent on node {curr_node.name}')
+
+        # Local import to avoid runtime circular dependencies with Context
+        from ..workflow.utils._transfer_utils import resolve_and_derive_transfer_context
+
+        target_agent, next_parent_ctx = resolve_and_derive_transfer_context(
+            target_name=target_name,
+            current_agent=curr_node,
+            root_agent=root_agent,
+            curr_ctx=child_ctx,
+            curr_parent_ctx=curr_parent_ctx,
         )
-      if child_ctx.interrupt_ids:
-        # Propagate child's interrupt_ids to this node's ctx
-        # so NodeRunner sees them after catching the error.
-        self._interrupt_ids.update(child_ctx.interrupt_ids)
-        raise NodeInterruptedError()
-      # When the caller passes raise_on_wait=True, surface a child
-      # that's WAITING (wait_for_output, no output, not transferring)
-      # as NodeInterruptedError so the parent's NodeRunner records
-      # the parent as WAITING instead of falsely COMPLETED.
-      if (
-          raise_on_wait
-          and built_node.wait_for_output
-          and child_ctx.output is None
-          and not child_ctx.actions.transfer_to_agent
-      ):
-        raise NodeInterruptedError()
-      return child_ctx if return_ctx else child_ctx.output
-
-    # Mode 2: Standalone execution (outside of workflow).
-    # Run the node directly via NodeRunner.
-    result = await self._run_node_standalone(
-        built_node,
-        node_input,
-        use_as_output=use_as_output,
-        use_sub_branch=use_sub_branch,
-        override_branch=override_branch,
-        override_isolation_scope=override_isolation_scope,
-        run_id=run_id,
-        resume_inputs=resume_inputs,
-    )
-    if result.error:
-      from ..workflow import _errors
+        if not target_agent:
+          raise ValueError(f"Transfer target agent '{target_name}' not found.")
+        if not next_parent_ctx:
+          available = []
+          if hasattr(curr_node, '_get_available_agent_names'):
+            available = curr_node._get_available_agent_names()
+          available_str = (
+              f"\nAvailable agents: {', '.join(available)}" if available else ''
+          )
+          raise ValueError(
+              f"Cannot transfer from '{curr_node.name}' to unrelated agent"
+              f" '{target_name}'.{available_str}"
+          )
+        curr_parent_ctx = next_parent_ctx
 
-      raise _errors.DynamicNodeFailError(
-          message=f'Dynamic node {built_node.name} failed',
-          error=result.error,
-          error_node_path=result.error_node_path,
-      )
-    if (
-        raise_on_wait
-        and built_node.wait_for_output
-        and result.output is None
-        and not result.actions.transfer_to_agent
-    ):
-      from ..workflow._errors import NodeInterruptedError
-
-      raise NodeInterruptedError()
-    return result if return_ctx else result.output
+        # Set up parameters for next iteration (the transfer target).
+        curr_node = target_agent
+        curr_run_id = None
+        curr_input = None  # Input for transfer target is usually empty.
+        resume_inputs = None
+
+        if not curr_parent_ctx:
+          raise AssertionError(
+              'curr_parent_ctx cannot be None during active workflow execution'
+          )
+
+        continue
+
+      # If no transfer occurred, execution of the branch is complete.
+      if return_ctx:
+        return child_ctx
+      return child_ctx.output
 
   # ============================================================================
   # Artifact methods

src/google/adk/runners.py

@@ -453,7 +453,6 @@ async def _run_node_async(
 
     Events flow through ic._event_queue via NodeRunner.
     """
-    from .workflow._node_runner import NodeRunner
 
     with tracer.start_as_current_span('invocation'):
       # 1. Setup
@@ -518,6 +517,8 @@ async def _run_node_async(
       from .agents.base_agent import BaseAgent
       from .agents.context import Context
       from .workflow._dynamic_node_scheduler import DynamicNodeScheduler
+      from .workflow._errors import DynamicNodeFailError
+      from .workflow._errors import NodeInterruptedError
       from .workflow._workflow import _LoopState
 
       root_ctx = Context(ic)
@@ -536,9 +537,6 @@ async def _run_node_async(
       # originating function-call id and so remain invisible to the
       # coordinator's view.
 
-      if not use_scheduler:
-        root_node_runner = NodeRunner(node=root_agent, parent_ctx=root_ctx)
-
       done_sentinel = object()
 
       async def _drive_root_node():
@@ -548,19 +546,18 @@ async def _drive_root_node():
             # Stateful live EUC/LRO streams may rehydrate freshly if not yet persisted.
             scheduler = DynamicNodeScheduler(state=_LoopState())
             root_ctx._workflow_scheduler = scheduler
-            ctx = await scheduler(
-                root_ctx,
+
+          try:
+            await root_ctx._run_node_internal(
                 root_agent,
-                node_input,
-                run_id='1',
-            )
-          else:
-            ctx = await root_node_runner.run(
                 node_input=node_input,
                 resume_inputs=resume_inputs,
             )
-          if ctx.error:
-            raise ctx.error
+          except NodeInterruptedError:
+            # The node was interrupted (e.g. for HITL).
+            pass
+          except DynamicNodeFailError as e:
+            raise e.error
         finally:
           await ic._event_queue.put((done_sentinel, None))
 
@@ -597,7 +594,8 @@ async def _run_node_live(
     """Run a non-agent BaseNode in live mode."""
     from .agents.context import Context
     from .workflow._dynamic_node_scheduler import DynamicNodeScheduler
-    from .workflow._node_runner import NodeRunner
+    from .workflow._errors import DynamicNodeFailError
+    from .workflow._errors import NodeInterruptedError
     from .workflow._workflow import _LoopState
     from .workflow._workflow import Workflow
 
@@ -612,28 +610,23 @@ async def _run_node_live(
     root_agent = self.agent
     is_workflow = isinstance(root_agent, Workflow)
 
-    if not is_workflow:
-      root_node_runner = NodeRunner(node=root_agent, parent_ctx=root_ctx)
-
     done_sentinel = object()
 
     async def _drive_root_node():
       try:
         if is_workflow:
           scheduler = DynamicNodeScheduler(state=_LoopState())
           root_ctx._workflow_scheduler = scheduler
-          ctx = await scheduler(
-              root_ctx,
+
+        try:
+          await root_ctx.run_node(
               root_agent,
-              None,
-              run_id='1',
-          )
-        else:
-          ctx = await root_node_runner.run(
               node_input=None,
           )
-        if ctx.error:
-          raise ctx.error
+        except NodeInterruptedError:
+          pass
+        except DynamicNodeFailError as e:
+          raise e.error
       finally:
         await ic._event_queue.put((done_sentinel, None))
 
@@ -999,9 +992,9 @@ async def run_async(
           agent_to_run = self.agent
         else:
           agent_to_run = self._find_agent_to_run(session, self.agent)
-        from .workflow.utils._workflow_graph_utils import build_node  # pylint: disable=g-import-not-at-top
 
-        agent_to_run = build_node(agent_to_run)
+        # The agent_to_run will be built/cloned inside Context.run_node,
+        # so we don't call build_node here to avoid double cloning.
       else:
         raise ValueError(
             "LlmAgent as root agent must have mode='chat', but got"