diff --git a/src/google/adk/workflow/_node_runner.py b/src/google/adk/workflow/_node_runner.py index 5ee004c3875..a6866f62df4 100644 --- a/src/google/adk/workflow/_node_runner.py +++ b/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) diff --git a/tests/unittests/workflow/test_dynamic_use_as_output.py b/tests/unittests/workflow/test_dynamic_use_as_output.py index ea8f67b6027..4986aa26a40 100644 --- a/tests/unittests/workflow/test_dynamic_use_as_output.py +++ b/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 diff --git a/tests/unittests/workflow/test_node_runner_ctx.py b/tests/unittests/workflow/test_node_runner_ctx.py index 953192b2da0..31811f5a49f 100644 --- a/tests/unittests/workflow/test_node_runner_ctx.py +++ b/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 --- diff --git a/tests/unittests/workflow/test_workflow_routes.py b/tests/unittests/workflow/test_workflow_routes.py index cfc125225c1..ef011ef1253 100644 --- a/tests/unittests/workflow/test_workflow_routes.py +++ b/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