feat(workflow): Use SingleAgentReactNode for leaf single_turn agents in new Workflow

DeanChensj · DeanChensj · commit dc2adf9869a8 · 2026-03-27T18:36:56.000-07:00
In the new _workflow_class.py Workflow, leaf single_turn LlmAgents now
run through SingleAgentReactNode instead of _SingleLlmAgent. This
enables proper event deduplication via _output_delegated — only the
LlmCallNode content event appears in the stream, not a duplicate
wrapper output event.

The old _workflow.py Workflow retains the original _SingleLlmAgent
path. Runtime detection uses event_queue presence on InvocationContext
to distinguish the two paths.

Also updates the use_as_output sample to use an LLM agent as the
delegated child node, and rewrites wrapper tests per adk-style
guidelines covering both old and new workflow paths.

Change-Id: I96aea3f12fe432d0449268eb75580c900fd047d7
diff --git a/contributing/workflow_samples/use_as_output/README.md b/contributing/workflow_samples/use_as_output/README.md
@@ -6,9 +6,11 @@ This sample demonstrates how to use `ctx.run_node(node, use_as_output=True)` to
 
 When `use_as_output=True` is set, the child node's output replaces the parent's output. The parent's own output event is suppressed to avoid duplication, and the child's output flows downstream through the graph as if the parent produced it.
 
+The child node can be any node type — this sample uses a single_turn LLM agent (`summarizer`) as the delegated child.
+
 ## Sample Inputs
 
-- Any text input (e.g. `hello world`)
+- Any text input (e.g. `The quick brown fox jumped over the lazy dog near the riverbank on a warm summer afternoon`)
 
 ## Graph
 
@@ -17,33 +19,38 @@ When `use_as_output=True` is set, the child node's output replaces the parent's
           |
           v
     [orchestrate]
-          |  (delegates output to [transform] via ctx.run_node)
+          |  (delegates output to [summarizer] via ctx.run_node)
           v
       [finalize]
 ```
 
 ## How To
 
-1. **Mark the orchestrator as rerun_on_resume**: The parent node that calls `ctx.run_node` must use `@node(rerun_on_resume=True)`.
+1. **Define the child node**: The child can be a function or an LLM agent. Its output becomes the parent's output and flows to downstream nodes.
+
+   ```python
+   from google.adk import Agent
+
+   summarizer = Agent(
+       name='summarizer',
+       model='gemini-2.5-flash',
+       instruction='Summarize the following text in one sentence.',
+   )
+   ```
+
+2. **Mark the orchestrator as rerun_on_resume**: The parent node that calls `ctx.run_node` must use `@node(rerun_on_resume=True)`.
 
    ```python
    from google.adk.workflow import node
 
    @node(rerun_on_resume=True)
    async def orchestrate(ctx: Context, node_input: str) -> str:
        return await ctx.run_node(
-           transform, node_input=node_input, use_as_output=True
+           summarizer, node_input=node_input, use_as_output=True
        )
    ```
 
-2. **Define the child node**: The child is a plain function. Its output becomes the parent's output and flows to downstream nodes.
-
-   ```python
-   def transform(node_input: str) -> str:
-       return node_input.upper()
-   ```
-
-3. **Downstream receives delegated output**: The `finalize` node receives the child's output (`"HELLO WORLD"`) as its `node_input`, not the parent's.
+3. **Downstream receives delegated output**: The `finalize` node receives the LLM's summary as its `node_input`, not the parent's.
 
    ```python
    def finalize(node_input: str) -> str:
diff --git a/contributing/workflow_samples/use_as_output/agent.py b/contributing/workflow_samples/use_as_output/agent.py
@@ -12,20 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.adk import Agent
 from google.adk import Context
 from google.adk.workflow import node
 from google.adk.workflow._base_node import START
 from google.adk.workflow._workflow_class import Workflow
 
 
-def transform(node_input: str) -> str:
-  return node_input.upper()
+summarizer = Agent(
+    name='summarizer',
+    model='gemini-2.5-flash',
+    instruction='Summarize the following text in one sentence.',
+)
 
 
 @node(rerun_on_resume=True)
 async def orchestrate(ctx: Context, node_input: str) -> str:
   return await ctx.run_node(
-      transform, node_input=node_input, use_as_output=True
+      summarizer, node_input=node_input, use_as_output=True
   )
 
 
diff --git a/src/google/adk/agents/llm/_single_agent_react_node.py b/src/google/adk/agents/llm/_single_agent_react_node.py
@@ -83,12 +83,12 @@ async def _run_impl(
       ctx: Context,
       node_input: Any,
   ) -> AsyncGenerator[Any, None]:
-    # Ensure IC has the agent — needed by downstream code
-    # (e.g. _build_response_event reads invocation_context.agent.name).
-    # When run via Runner(node=...), ic.agent is None.
+    # Set IC agent — needed by downstream code
+    # (e.g. _build_basic_request reads agent.canonical_model).
+    # When run inside a Workflow, ic.agent may be the Workflow itself
+    # which lacks canonical_model. Always override.
     # TODO: remove this dependency.
-    if ctx._invocation_context.agent is None:
-      ctx._invocation_context.agent = self.agent
+    ctx._invocation_context.agent = self.agent
 
     # Always provide our own scheduler so child nodes
     # (LlmCallNode, ParallelToolCallNode) are managed by this node,
diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py
@@ -16,6 +16,8 @@
 
 - Sets a branch for content isolation (single_turn mode only)
 - Converts node_input to user content (single_turn mode only)
+- Runs leaf single_turn agents via SingleAgentReactNode (new Workflow)
+  or _SingleLlmAgent (old Workflow)
 - Re-emits finish_task output so the outer node_runner can route it
 """
 
@@ -57,8 +59,12 @@ class _LlmAgentWrapper(BaseNode):
   """Adapts a task/single_turn LlmAgent for use as a workflow graph node.
 
   Output handling by mode:
-    single_turn (leaf, no sub_agents): Bypasses Mesh by running
-      _SingleLlmAgent directly. Output is extracted via
+    single_turn (leaf, no sub_agents, new Workflow): Runs the ReAct
+      loop via SingleAgentReactNode. LlmCallNode events are enqueued
+      internally; the wrapper post-processes the final text output
+      (output_schema validation, output_key storage).
+    single_turn (leaf, no sub_agents, old Workflow): Bypasses Mesh by
+      running _SingleLlmAgent directly. Output is extracted via
       LlmAgent._maybe_save_output_to_state and emitted as a
       separate Event before END_OF_AGENT.
     single_turn (with sub_agents): Runs the full LlmAgent which
@@ -72,6 +78,7 @@ class _LlmAgentWrapper(BaseNode):
   agent: LlmAgent = Field(...)
   rerun_on_resume: bool = Field(default=True)
   _single: Any = PrivateAttr(default=None)
+  _react: Any = PrivateAttr(default=None)
 
   @model_validator(mode='before')
   @classmethod
@@ -96,12 +103,21 @@ def _validate_and_default_mode(self) -> _LlmAgentWrapper:
     if self.agent.mode == 'task':
       self.wait_for_output = True
 
-    # For leaf single_turn agents, use _SingleLlmAgent directly,
-    # bypassing the _Mesh orchestration layer.
+    # For leaf single_turn agents, prepare both execution paths.
+    # The old Workflow (uses _node_runner.py, no event_queue) needs
+    # _SingleLlmAgent. The new Workflow (_workflow_class.py, sets
+    # event_queue) uses SingleAgentReactNode. The choice is made at
+    # runtime in _run_impl based on event_queue presence.
     if self.agent.mode == 'single_turn' and not self.agent.sub_agents:
+      from ..agents.llm._single_agent_react_node import (
+          SingleAgentReactNode,
+      )
       from ..agents.llm._single_llm_agent import _SingleLlmAgent
 
       self._single = _SingleLlmAgent.from_base_llm_agent(self.agent)
+      self._react = SingleAgentReactNode(
+          name=self.agent.name, agent=self.agent
+      )
 
     return self
 
@@ -122,6 +138,10 @@ def model_copy(
         copied._single = copied._single.model_copy(
             update={'name': update['name']}
         )
+      if copied._react is not None:
+        copied._react = copied._react.model_copy(
+            update={'name': update['name'], 'agent': copied.agent}
+        )
     return copied
 
   def _validate_input(self, node_input: Any) -> None:
@@ -153,13 +173,28 @@ def _prepare_input(
     ic = ctx._invocation_context.model_copy(
         update={'branch': branch},
     )
+    # event_queue is excluded from model_copy; propagate manually
+    # so child NodeRunners can enqueue events.
+    ic.event_queue = ctx._invocation_context.event_queue
     agent_ctx = Context(
         invocation_context=ic,
         node_path=ctx.node_path,
         execution_id=ctx.execution_id,
     )
     return agent_ctx, agent_input
 
+  def _use_react_path(self, ctx: Context) -> bool:
+    """Returns True if we should use SingleAgentReactNode (new Workflow).
+
+    The new Workflow (_workflow_class.py) runs via Runner._run_node_async
+    which sets ic.event_queue. The old Workflow (_workflow.py) runs via
+    Runner.run_async (agent path) which does not set event_queue.
+    """
+    return (
+        self._react is not None
+        and ctx._invocation_context.event_queue is not None
+    )
+
   @override
   async def _run_impl(
       self,
@@ -171,6 +206,47 @@ async def _run_impl(
     self._validate_input(node_input)
     agent_ctx, agent_input = self._prepare_input(ctx, node_input)
 
+    if self._use_react_path(ctx):
+      # New Workflow: leaf single_turn via SingleAgentReactNode.
+      # Inject input as user content in session, then run the react
+      # loop. LlmCallNode events are enqueued to event_queue internally;
+      # only the final text output comes through the generator.
+      if agent_input is not None:
+        ic = agent_ctx._invocation_context
+        ic.session.events.append(
+            Event(
+                invocation_id=ic.invocation_id,
+                author='user',
+                branch=ic.branch,
+                content=agent_input,
+            )
+        )
+      output = None
+      async for event in self._react.run(
+          ctx=agent_ctx, node_input=None
+      ):
+        if isinstance(event, Event) and event.output is not None:
+          output = event.output
+
+      if output is not None:
+        if isinstance(output, str) and self.agent.output_schema:
+          if not output.strip():
+            return
+          from ..utils._schema_utils import validate_schema
+
+          output = validate_schema(self.agent.output_schema, output)
+        if self.agent.output_key:
+          ctx.actions.state_delta[self.agent.output_key] = output
+        # Mark output as delegated so the wrapper's NodeRunner
+        # captures the value without enqueuing a separate output
+        # event. The LlmCallNode content event (already enqueued)
+        # carries the visible response; this avoids duplication.
+        ctx._output_delegated = True
+        yield output
+      return
+      yield  # noqa: unreachable — keeps this an async generator
+
+    # Determine inner runner: _SingleLlmAgent for leaf, agent for others.
     inner = self._single if self._single is not None else self.agent
 
     # When the agent has parallel_worker=True, call run_node_impl()
@@ -182,8 +258,8 @@ async def _run_impl(
 
     if self.agent.mode == 'single_turn':
       if self._single is not None:
-        # Leaf agent bypass: since we skip LlmAgent.run_node_impl(),
-        # replicate its output handling here.
+        # Old Workflow: leaf agent bypass. Since we skip
+        # LlmAgent.run_node_impl(), replicate its output handling.
         # _maybe_save_output_to_state applies output_schema/output_key
         # and clears content on the final response. We emit the output
         # as a pathless Event before END_OF_AGENT.
diff --git a/tests/unittests/workflow/test_llm_agent_wrapper.py b/tests/unittests/workflow/test_llm_agent_wrapper.py
