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 <deanchen@google.com>
PiperOrigin-RevId: 949367240

Author: DeanChensj

src/google/adk/agents/context.py

@@ -613,18 +613,18 @@ async def _run_node_internal(
           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.

src/google/adk/tools/_node_tool.py

@@ -142,13 +142,16 @@ async def run_async(
     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

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:

tests/unittests/workflow/test_node_tool.py

@@ -328,6 +328,53 @@ def greet_node(request: str) -> str:
   ].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."""

tests/unittests/workflow/test_workflow_nested.py

@@ -27,9 +27,10 @@
 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 _run_impl(
     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)