From fd006db9153fe6c51f281f36e35b40d243e5fca0 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Thu, 16 Jul 2026 22:26:38 -0700 Subject: [PATCH] fix(workflow): fix task agent resumption in nested workflows Fix an issue where task-mode LLM agents invoked inside nested workflows or custom nodes with raise_on_wait=True fail to resume upon user reply. - Ensure raise_on_wait=True in Context._run_node_internal triggers NodeInterruptedError when child output is None regardless of wait_for_output defaults. - Update ReplayInterceptor Case 5 to allow re-running Workflow instances and rerun_on_resume nodes when recovered output is None. - Add unit test test_nested_workflow_with_task_agent to cover task agent resumption in nested workflows. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 949367240 --- src/google/adk/agents/context.py | 18 +++---- src/google/adk/tools/_node_tool.py | 5 +- .../adk/workflow/utils/_replay_interceptor.py | 11 +++- tests/unittests/workflow/test_node_tool.py | 47 +++++++++++++++++ .../workflow/test_workflow_nested.py | 52 ++++++++++++++++++- 5 files changed, 120 insertions(+), 13 deletions(-) diff --git a/src/google/adk/agents/context.py b/src/google/adk/agents/context.py index 1679da8e..da256003 100644 --- a/src/google/adk/agents/context.py +++ b/src/google/adk/agents/context.py @@ -613,18 +613,18 @@ class Context(ReadonlyContext): 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) + # execution 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 + if raise_on_wait and child_ctx.output is None and not transfer_to_agent: + from ..workflow._workflow import Workflow - raise NodeInterruptedError() + if isinstance(curr_node, Workflow) or getattr( + curr_node, 'wait_for_output', False + ): + 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. diff --git a/src/google/adk/tools/_node_tool.py b/src/google/adk/tools/_node_tool.py index f69c0ef3..509d4dc4 100644 --- a/src/google/adk/tools/_node_tool.py +++ b/src/google/adk/tools/_node_tool.py @@ -142,13 +142,16 @@ class NodeTool(BaseTool): tool_branch = f'{base_branch}.{segment}' if base_branch else segment try: - return await tool_context.run_node( + res = await tool_context.run_node( self.node, node_input=node_input, override_branch=tool_branch, use_sub_branch=False, raise_on_wait=True, ) + if res is None: + return {'result': None} + return res except NodeInterruptedError as nie: # Propagates the interrupt up so the runner pauses the invocation raise nie diff --git a/src/google/adk/workflow/utils/_replay_interceptor.py b/src/google/adk/workflow/utils/_replay_interceptor.py index b8780fa5..23f130af 100644 --- a/src/google/adk/workflow/utils/_replay_interceptor.py +++ b/src/google/adk/workflow/utils/_replay_interceptor.py @@ -61,6 +61,7 @@ def check_interception( current_run: DynamicNodeRun | None = None, ) -> InterceptionResult: """Determine if a node execution should be intercepted based on history.""" + from .._workflow import Workflow # pylint: disable=g-import-not-at-top # Case 1: Same-turn completed or waiting interception (dynamic nodes only). # If a node already successfully executed or is currently blocked in the @@ -127,8 +128,14 @@ def check_interception( else: # Case 5: Cross-turn no events, or events contain no output, route, or interrupts. - # Rerun Workflow nodes to guide nested children; otherwise fall through. - if getattr(node, "wait_for_output", False) and recovered.output is None: + # Rerun Workflow nodes, wait_for_output nodes, and rerun_on_resume nodes + # with no prior output so they can guide nested children or resume execution; + # otherwise fall through. + if ( + isinstance(node, Workflow) + or getattr(node, "wait_for_output", False) + or getattr(node, "rerun_on_resume", False) + ) and recovered.output is None: should_run = True resume_inputs = recovered.resolved_responses else: diff --git a/tests/unittests/workflow/test_node_tool.py b/tests/unittests/workflow/test_node_tool.py index cdbbc891..d0f7de0e 100644 --- a/tests/unittests/workflow/test_node_tool.py +++ b/tests/unittests/workflow/test_node_tool.py @@ -328,6 +328,53 @@ async def test_function_node_wrapped_as_tool_returns_output( ].function_response.response == {'result': 'Hello, world!'} +@pytest.mark.asyncio +async def test_function_node_wrapped_as_tool_no_output( + request: pytest.FixtureRequest, +): + """NodeTool wrapping a function node that returns None completes with None result.""" + + @node + def no_output_node(request: str): + yield Event(output=None) + + no_output_node.input_schema = GreetRequest + no_output_tool = NodeTool(node=no_output_node, name='no_output_tool') + + parent_agent = LlmAgent( + name='parent_agent', + model=testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call( + name='no_output_tool', + args={'request': 'world'}, + ), + types.Part.from_text(text='Processed no output.'), + ] + ), + tools=[no_output_tool], + ) + + app = App( + name=request.function.__name__, + root_agent=parent_agent, + ) + runner = testing_utils.InMemoryRunner(app=app) + events = await runner.run_async( + testing_utils.get_user_content('Run no output') + ) + + func_response_events = [ + e + for e in events + if e.content and e.content.parts and e.content.parts[0].function_response + ] + assert len(func_response_events) == 1 + assert func_response_events[0].content.parts[ + 0 + ].function_response.response == {'result': None} + + @pytest.mark.asyncio async def test_workflow_tool_with_join_node(request: pytest.FixtureRequest): """WorkflowTool containing a JoinNode works correctly when wrapped as a tool.""" diff --git a/tests/unittests/workflow/test_workflow_nested.py b/tests/unittests/workflow/test_workflow_nested.py index 90172006..34df5766 100644 --- a/tests/unittests/workflow/test_workflow_nested.py +++ b/tests/unittests/workflow/test_workflow_nested.py @@ -27,9 +27,10 @@ from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.long_running_tool import LongRunningFunctionTool from google.adk.workflow import BaseNode from google.adk.workflow import JoinNode +from google.adk.workflow import node +from google.adk.workflow import Workflow from google.adk.workflow._base_node import START from google.adk.workflow._node_status import NodeStatus -from google.adk.workflow._workflow import Workflow from google.adk.workflow.utils._workflow_hitl_utils import create_request_input_response from google.adk.workflow.utils._workflow_hitl_utils import get_request_input_interrupt_ids from google.adk.workflow.utils._workflow_hitl_utils import has_request_input_function_call @@ -1122,3 +1123,52 @@ async def test_nested_workflow_partial_resume(): if e.long_running_tool_ids: final_interrupts.update(e.long_running_tool_ids) assert not final_interrupts + + +@pytest.mark.asyncio +async def test_nested_workflow_with_task_agent(request: pytest.FixtureRequest): + """Tests that a task-mode LlmAgent inside a nested Workflow re-runs on user reply.""" + mock_model = testing_utils.MockModel.create( + responses=[ + types.Part.from_text(text='Please provide the secret code:'), + types.Part.from_function_call( + name='finish_task', + args={'result': 'Success with secret'}, + ), + ] + ) + task_agent = LlmAgent( + name='inner_task_agent', + model=mock_model, + mode='task', + ) + inner_wf = Workflow( + name='inner_wf', + edges=[('START', task_agent)], + ) + + @node(rerun_on_resume=True) + async def outer_driver(ctx: Context, node_input: Any): + res = await ctx.run_node(inner_wf, node_input='start', raise_on_wait=True) + yield Event(output=f'outer: {res}') + + outer_wf = Workflow( + name='outer_wf', + edges=[('START', outer_driver)], + ) + + app = App(name=request.function.__name__, root_agent=outer_wf) + runner = testing_utils.InMemoryRunner(app=app) + + events1 = await runner.run_async('hello') + texts1 = [ + p.text + for e in events1 + if e.content and e.content.parts + for p in e.content.parts + if p.text + ] + assert 'Please provide the secret code:' in texts1 + + events2 = await runner.run_async('secret_code_123') + assert any('outer: ' in str(e.output) for e in events2)