fix(workflow): preserve routes when flushing and delegating output

src/google/adk/workflow/_node_runner.py

@@ -45,7 +45,11 @@
 
 def _has_non_output_content(event: Event) -> bool:
   if event.actions:
-    if event.actions.state_delta or event.actions.artifact_delta:
+    if (
+        event.actions.state_delta
+        or event.actions.artifact_delta
+        or event.actions.route is not None
+    ):
       return True
   return False
 
@@ -492,7 +496,6 @@ async def _enqueue_event(self, event: Event, ctx: Context) -> None:
   async def _flush_output_and_deltas(self, ctx: Context) -> None:
     """Emit deferred output and/or unflushed state/artifact deltas."""
     from ..events.event import Event
-    from ..events.event_actions import EventActions
 
     state_delta = ctx.actions.state_delta
     artifact_delta = ctx.actions.artifact_delta
@@ -514,13 +517,7 @@ async def _flush_output_and_deltas(self, ctx: Context) -> None:
         output=ctx._output_value if has_deferred_output else None,
         route=ctx._route_value if has_unflushed_route else None,
     )
-    if has_deltas:
-      event.actions = EventActions(
-          state_delta=dict(state_delta),
-          artifact_delta=dict(artifact_delta),
-      )
-      state_delta.clear()
-      artifact_delta.clear()
+    self._flush_deltas(event, ctx)
 
     self._enrich_event(event, ctx)
     await ctx._invocation_context._enqueue_event(event)

tests/unittests/workflow/test_dynamic_use_as_output.py

@@ -50,6 +50,36 @@ def _make_app(name: str, agent: Workflow, resumable: bool) -> App:
   )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('route', ['next', False, 0, []])
+async def test_delegated_output_preserves_parent_route(
+    request: pytest.FixtureRequest, route: bool | int | str | list[str]
+) -> None:
+  """Suppressing a duplicate output keeps its parent's route event."""
+
+  def child() -> str:
+    return 'from_child'
+
+  child_node = FunctionNode(func=child)
+
+  async def parent(ctx: Context) -> Event:
+    output = await ctx.run_node(child_node, use_as_output=True)
+    return Event(output=output, route=route)
+
+  parent_node = FunctionNode(func=parent, rerun_on_resume=True)
+  wf = Workflow(name='delegated_route', edges=[(START, parent_node)])
+  runner = testing_utils.InMemoryRunner(
+      app=App(name=request.function.__name__, root_agent=wf)
+  )
+
+  events = await runner.run_async(testing_utils.get_user_content('start'))
+
+  assert _get_outputs(events) == ['from_child']
+  assert [
+      event.actions.route for event in events if event.actions.route is not None
+  ] == [route]
+
+
 # ---------------------------------------------------------------------------
 
 # FunctionNode delegates to FunctionNode

tests/unittests/workflow/test_node_runner_ctx.py

@@ -20,6 +20,8 @@
 forward correctly.
 """
 
+from collections.abc import AsyncGenerator
+from typing import Any
 from unittest import mock
 from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
@@ -275,6 +277,41 @@ async def _run_impl(self, *, ctx, node_input):
   assert child_ctx.route == 'branch_a'
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('with_state', [False, True])
+@pytest.mark.parametrize('with_artifact', [False, True])
+async def test_deferred_route_is_emitted_with_pending_deltas(
+    with_state: bool, with_artifact: bool
+) -> None:
+  """A final event keeps its route alongside state and artifact updates."""
+
+  class _Node(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Event, None]:
+      ctx.output = 'result'
+      ctx.route = 'branch_a'
+      if with_state:
+        ctx.state['visited'] = True
+      if with_artifact:
+        ctx.actions.artifact_delta['report.txt'] = 1
+      return
+      yield
+
+  parent_ctx, events = _make_ctx()
+  await NodeRunner(node=_Node(name='router'), parent_ctx=parent_ctx).run()
+
+  assert len(events) == 1
+  event = events[0]
+  assert event.output == 'result'
+  assert event.actions.route == 'branch_a'
+  assert event.actions.state_delta == ({'visited': True} if with_state else {})
+  assert event.actions.artifact_delta == (
+      {'report.txt': 1} if with_artifact else {}
+  )
+
+
 # --- ctx.interrupt_ids ---
 
 

tests/unittests/workflow/test_workflow_routes.py

@@ -14,11 +14,15 @@
 
 """Testings for the Workflow routes."""
 
+from collections.abc import AsyncGenerator
 from typing import Any
 
 from google.adk.agents.context import Context
 from google.adk.apps.app import App
+from google.adk.events.event import Event
+from google.adk.workflow import BaseNode
 from google.adk.workflow import Edge
+from google.adk.workflow import FunctionNode
 from google.adk.workflow import node
 from google.adk.workflow import START
 from google.adk.workflow._graph import DEFAULT_ROUTE
@@ -776,9 +780,11 @@ async def test_fan_in_out_with_route(request: pytest.FixtureRequest):
     await runner.run_async(testing_utils.get_user_content('start'))
 
 
+@pytest.mark.parametrize('router_kind', ['event', 'deferred', 'delegated'])
 @pytest.mark.asyncio
 async def test_parent_retry_reruns_sub_workflow_along_recorded_route(
     request: pytest.FixtureRequest,
+    router_kind: str,
 ) -> None:
   """A parent's retry reruns a sub-workflow along its recorded route."""
   runs = {'router': 0, 'flaky_b': 0, 'node_c': 0}
@@ -797,7 +803,34 @@ def node_c(node_input: Any) -> str:
     runs['node_c'] += 1
     return 'c_done'
 
-  router = TestingNode(name='router', output='routed', route=pick_route)
+  class DeferredRouter(BaseNode):
+
+    async def _run_impl(
+        self, *, ctx: Context, node_input: Any
+    ) -> AsyncGenerator[Event, None]:
+      ctx.output = 'routed'
+      ctx.route = pick_route(ctx, node_input)
+      ctx.state['route_visits'] = runs['router']
+      return
+      yield
+
+  def output_source() -> str:
+    return 'routed'
+
+  child = FunctionNode(func=output_source)
+
+  async def delegated_router(ctx: Context) -> Event:
+    output = await ctx.run_node(child, use_as_output=True)
+    return Event(output=output, route=pick_route(ctx, None))
+
+  if router_kind == 'deferred':
+    router = DeferredRouter(name='router')
+  elif router_kind == 'delegated':
+    router = FunctionNode(
+        name='router', func=delegated_router, rerun_on_resume=True
+    )
+  else:
+    router = TestingNode(name='router', output='routed', route=pick_route)
   inner_wf = Workflow(
       name='inner_wf',
       edges=[(START, router), (router, {'go_b': flaky_b, 'go_c': node_c})],
@@ -818,3 +851,7 @@ async def driver(ctx: Context) -> Any:
 
   assert runs == {'router': 1, 'flaky_b': 2, 'node_c': 0}
   assert 'b_done' in [e.output for e in events if e.output is not None]
+  if router_kind == 'deferred':
+    assert runner.session.state['route_visits'] == 1
+  if router_kind == 'delegated':
+    assert sum(event.output == 'routed' for event in events) == 1