fix: Yield the long running tool response before pausing execution

PiperOrigin-RevId: 825056377

src/google/adk/agents/llm_agent.py

@@ -430,16 +430,24 @@ class LlmAgent(BaseAgent):
       yield self._create_agent_state_event(ctx)
       return
 
+    should_pause = False
     async with Aclosing(self._llm_flow.run_async(ctx)) as agen:
       async for event in agen:
         self.__maybe_save_output_to_state(event)
         yield event
         if ctx.should_pause_invocation(event):
-          return
+          # Do not pause immediately, wait until the long running tool call is
+          # executed.
+          should_pause = True
+    if should_pause:
+      return
 
     if ctx.is_resumable:
       events = ctx._get_events(current_invocation=True, current_branch=True)
-      if events and ctx.should_pause_invocation(events[-1]):
+      if events and (
+          ctx.should_pause_invocation(events[-1])
+          or ctx.should_pause_invocation(events[-2])
+      ):
         return
       # Only yield an end state if the last event is no longer a long running
       # tool call.

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -383,16 +383,30 @@ class BaseLlmFlow(ABC):
     events = invocation_context._get_events(
         current_invocation=True, current_branch=True
     )
+
+    # Long running tool calls should have been handled before this point.
+    # If there are still long running tool calls, it means the agent is paused
+    # before, and its branch hasn't been resumed yet.
+    if (
+        invocation_context.is_resumable
+        and events
+        and len(events) > 1
+        # TODO: here we are using the last 2 events to decide whether to pause
+        # the invocation. But this is just being optmisitic, we should find a
+        # way to pause when the long running tool call is followed by more than
+        # one text responses.
+        and (
+            invocation_context.should_pause_invocation(events[-1])
+            or invocation_context.should_pause_invocation(events[-2])
+        )
+    ):
+      return
+
     if (
         invocation_context.is_resumable
         and events
         and events[-1].get_function_calls()
     ):
-      # Long running tool calls should have been handled before this point.
-      # If there are still long running tool calls, it means the agent is paused
-      # before, and its branch hasn't been resumed yet.
-      if invocation_context.should_pause_invocation(events[-1]):
-        return
       model_response_event = events[-1]
       async with Aclosing(
           self._postprocess_handle_function_calls_async(

src/google/adk/tools/function_tool.py

@@ -205,6 +205,7 @@ You could retry calling this tool, but it is IMPORTANT for you to provide all th
                 ' ToolConfirmation payload.'
             ),
         )
+        tool_context.actions.skip_summarization = True
         return {
             'error': (
                 'This tool call requires confirmation, please approve or'

tests/unittests/runners/test_pause_invocation.py

@@ -44,7 +44,7 @@ def _transfer_call_part(agent_name: str) -> Part:
 
 
 def test_tool() -> str:
-  return ""
+  return "result"
 
 
 class _TestingAgent(BaseAgent):
@@ -126,6 +126,12 @@ class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
     """Tests that a single LlmAgent pauses on long running function call."""
     assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
         ("root_agent", Part.from_function_call(name="test_tool", args={})),
+        (
+            "root_agent",
+            Part.from_function_response(
+                name="test_tool", response={"result": "result"}
+            ),
+        ),
     ]
 
 
@@ -168,6 +174,12 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
             ),
         ),
         ("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
+        (
+            "sub_agent_1",
+            Part.from_function_response(
+                name="test_tool", response={"result": "result"}
+            ),
+        ),
     ]
 
   @pytest.mark.asyncio
@@ -195,7 +207,7 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
         (
             "sub_agent_1",
             Part.from_function_response(
-                name="test_tool", response={"result": ""}
+                name="test_tool", response={"result": "result"}
             ),
         ),
         ("sub_agent_1", "model response after tool call"),
@@ -207,6 +219,12 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
             ),
         ),
         ("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
+        (
+            "sub_agent_2",
+            Part.from_function_response(
+                name="test_tool", response={"result": "result"}
+            ),
+        ),
     ]
 
 
@@ -384,6 +402,12 @@ class TestPauseInvocationWithLoopAgent(BasePauseInvocationTest):
             ),
         ),
         ("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
+        (
+            "sub_agent_2",
+            Part.from_function_response(
+                name="test_tool", response={"result": "result"}
+            ),
+        ),
     ]
 
 
@@ -435,6 +459,12 @@ class TestPauseInvocationWithLlmAgentTree(BasePauseInvocationTest):
         ("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
         ("sub_llm_agent_1", _TRANSFER_RESPONSE_PART),
         ("sub_llm_agent_2", Part.from_function_call(name="test_tool", args={})),
+        (
+            "sub_llm_agent_2",
+            Part.from_function_response(
+                name="test_tool", response={"result": "result"}
+            ),
+        ),
     ]
 
 
@@ -489,4 +519,10 @@ class TestPauseInvocationWithWithTransferLoop(BasePauseInvocationTest):
         ("sub_llm_agent_2", _transfer_call_part("root_agent")),
         ("sub_llm_agent_2", _TRANSFER_RESPONSE_PART),
         ("root_agent", Part.from_function_call(name="test_tool", args={})),
+        (
+            "root_agent",
+            Part.from_function_response(
+                name="test_tool", response={"result": "result"}
+            ),
+        ),
     ]

tests/unittests/runners/test_run_tool_confirmation.py

@@ -36,6 +36,16 @@ import pytest
 
 from .. import testing_utils
 
+HINT_TEXT = (
+    "Please approve or reject the tool call _test_function() by"
+    " responding with a FunctionResponse with an"
+    " expected ToolConfirmation payload."
+)
+
+TOOL_CALL_ERROR_RESPONSE = {
+    "error": "This tool call requires confirmation, please approve or reject."
+}
+
 
 def _create_llm_response_from_tools(
     tools: list[FunctionTool],
@@ -57,13 +67,9 @@ def _create_llm_response_from_text(text: str) -> GenerateContentResponse:
   )
 
 
-def _test_request_confirmation_function(
+def _test_function(
     tool_context: ToolContext,
 ) -> dict[str, str]:
-  """A test tool function that requests confirmation."""
-  if not tool_context.tool_confirmation:
-    tool_context.request_confirmation(hint="test hint for request_confirmation")
-    return {"error": "test error for request_confirmation"}
   return {"result": f"confirmed={tool_context.tool_confirmation.confirmed}"}
 
 
@@ -82,7 +88,7 @@ def _test_request_confirmation_function_with_custom_schema(
             }
         },
     )
-    return {"error": "test error for request_confirmation"}
+    return TOOL_CALL_ERROR_RESPONSE
   return {
       "result": f"confirmed={tool_context.tool_confirmation.confirmed}",
       "custom_payload": tool_context.tool_confirmation.payload,
@@ -104,7 +110,7 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
   @pytest.fixture
   def tools(self) -> list[FunctionTool]:
     """Provides the tools for the agent."""
-    return [FunctionTool(func=_test_request_confirmation_function)]
+    return [FunctionTool(func=_test_function, require_confirmation=True)]
 
   @pytest.fixture
   def llm_responses(
@@ -114,9 +120,6 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
     return [
         _create_llm_response_from_tools(tools),
         _create_llm_response_from_text("test llm response after tool call"),
-        _create_llm_response_from_text(
-            "test llm response after final tool call"
-        ),
     ]
 
   @pytest.fixture
@@ -163,7 +166,7 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
                             "args": {},
                         },
                         "toolConfirmation": {
-                            "hint": "test hint for request_confirmation",
+                            "hint": HINT_TEXT,
                             "confirmed": False,
                         },
                     },
@@ -174,12 +177,10 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
             agent.name,
             Part(
                 function_response=FunctionResponse(
-                    name=tools[0].name,
+                    name=tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
-                    response={"error": "test error for request_confirmation"},
                 )
             ),
         ),
-        (agent.name, "test llm response after tool call"),
     ]
 
     simplified = testing_utils.simplify_events(copy.deepcopy(events))
@@ -208,11 +209,13 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
             Part(
                 function_response=FunctionResponse(
                     name=tools[0].name,
-                    response={"result": f"confirmed={tool_call_confirmed}"},
+                    response={"result": f"confirmed={tool_call_confirmed}"}
+                    if tool_call_confirmed
+                    else {"error": "This tool call is rejected."},
                 )
             ),
         ),
-        (agent.name, "test llm response after final tool call"),
+        (agent.name, "test llm response after tool call"),
     ]
     for event in events:
       assert event.invocation_id != invocation_id
@@ -312,8 +315,7 @@ class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest):
             agent.name,
             Part(
                 function_response=FunctionResponse(
-                    name=tools[0].name,
+                    name=tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
-                    response={"error": "test error for request_confirmation"},
                 )
             ),
         ),
@@ -380,7 +382,7 @@ class TestHITLConfirmationFlowWithResumableApp:
   @pytest.fixture
   def tools(self) -> list[FunctionTool]:
     """Provides the tools for the agent."""
-    return [FunctionTool(func=_test_request_confirmation_function)]
+    return [FunctionTool(func=_test_function, require_confirmation=True)]
 
   @pytest.fixture
   def llm_responses(
@@ -409,8 +411,8 @@ class TestHITLConfirmationFlowWithResumableApp:
   @pytest.fixture
   def runner(self, agent: LlmAgent) -> testing_utils.InMemoryRunner:
     """Provides an in-memory runner for the agent."""
-    # Mark the app as resumable. So that the invocation will be paused after the
+    # Mark the app as resumable. So that the invocation will be paused when
-    # long running tool call.
+    # tool confirmation is requested.
     app = App(
         name="test_app",
         resumability_config=ResumabilityConfig(is_resumable=True),
@@ -427,8 +429,8 @@ class TestHITLConfirmationFlowWithResumableApp:
     """Tests HITL flow where all tool calls are confirmed."""
     events = runner.run("test user query")
 
-    # Verify that the invocation is paused after the long running tool call.
+    # Verify that the invocation is paused when tool confirmation is requested.
-    # So that no intermediate function response and llm response is generated.
+    # The tool call returns error response, and summarization was skipped.
     assert testing_utils.simplify_resumable_app_events(
         copy.deepcopy(events)
     ) == [
@@ -448,13 +450,21 @@ class TestHITLConfirmationFlowWithResumableApp:
                             "args": {},
                         },
                         "toolConfirmation": {
-                            "hint": "test hint for request_confirmation",
+                            "hint": HINT_TEXT,
                             "confirmed": False,
                         },
                     },
                 )
             ),
         ),
+        (
+            agent.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=agent.tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
+                )
+            ),
+        ),
     ]
     ask_for_confirmation_function_call_id = (
         events[1].content.parts[0].function_call.id
@@ -499,7 +509,7 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
   @pytest.fixture
   def tools(self) -> list[FunctionTool]:
     """Provides the tools for the agent."""
-    return [FunctionTool(func=_test_request_confirmation_function)]
+    return [FunctionTool(func=_test_function, require_confirmation=True)]
 
   @pytest.fixture
   def llm_responses(
@@ -535,8 +545,8 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
   @pytest.fixture
   def runner(self, agent: SequentialAgent) -> testing_utils.InMemoryRunner:
     """Provides an in-memory runner for the agent."""
-    # Mark the app as resumable. So that the invocation will be paused after the
+    # Mark the app as resumable. So that the invocation will be paused when
-    # long running tool call.
+    # tool confirmation is requested.
     app = App(
         name="test_app",
         resumability_config=ResumabilityConfig(is_resumable=True),
@@ -558,8 +568,8 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
     #   - sub_agent1 has a tool call that asks for HITL confirmation.
     #   - sub_agent2 does not have any tool calls.
     # - The test will:
-    #   - Run the query and verify that the invocation is paused after the long
+    #   - Run the query and verify that the invocation is paused when tool
-    #     running tool call, at sub_agent1.
+    #     confirmation is requested, at sub_agent1.
     #   - Resume the invocation and execute the tool call from sub_agent1.
     #   - Verify that root_agent continues to run sub_agent2.
 
@@ -568,8 +578,8 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
     sub_agent2 = agent.sub_agents[1]
 
     # Step 1:
-    # Verify that the invocation is paused after the long running tool call.
+    # Verify that the invocation is paused when tool confirmation is requested.
-    # So that no intermediate function response and llm response is generated.
+    # So that no intermediate llm response is generated.
     # And the second sub agent is not started.
     assert testing_utils.simplify_resumable_app_events(
         copy.deepcopy(events)
@@ -600,13 +610,22 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
                             "args": {},
                         },
                         "toolConfirmation": {
-                            "hint": "test hint for request_confirmation",
+                            "hint": HINT_TEXT,
                             "confirmed": False,
                         },
                     },
                 )
             ),
         ),
+        (
+            sub_agent1.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=sub_agent1.tools[0].name,
+                    response=TOOL_CALL_ERROR_RESPONSE,
+                )
+            ),
+        ),
     ]
     ask_for_confirmation_function_call_id = (
         events[2].content.parts[0].function_call.id
@@ -664,7 +683,7 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
   @pytest.fixture
   def tools(self) -> list[FunctionTool]:
     """Provides the tools for the agent."""
-    return [FunctionTool(func=_test_request_confirmation_function)]
+    return [FunctionTool(func=_test_function, require_confirmation=True)]
 
   @pytest.fixture
   def llm_responses(
@@ -702,8 +721,8 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
   @pytest.fixture
   def runner(self, agent: ParallelAgent) -> testing_utils.InMemoryRunner:
     """Provides an in-memory runner for the agent."""
-    # Mark the app as resumable. So that the invocation will be paused after the
+    # Mark the app as resumable. So that the invocation will be paused when
-    # long running tool call.
+    # tool confirmation is requested.
     app = App(
         name="test_app",
         resumability_config=ResumabilityConfig(is_resumable=True),
@@ -725,15 +744,15 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
     #   sub_agent2.
     # - Both sub_agents have a tool call that asks for HITL confirmation.
     # - The test will:
-    #   - Run the query and verify that each branch is paused after the long
+    #   - Run the query and verify that each branch is paused when tool
-    #     running tool call.
+    #     confirmation is requested.
     #   - Resume the invocation and execute the tool call of each branch.
 
     sub_agent1 = agent.sub_agents[0]
     sub_agent2 = agent.sub_agents[1]
 
     # Verify that each branch is paused after the long running tool call.
-    # So that no intermediate function response and llm response is generated.
+    # So that no intermediate llm response is generated.
     root_agent_events = [event for event in events if event.branch is None]
     sub_agent1_branch_events = [
         event
@@ -776,13 +795,22 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
                             "args": {},
                         },
                         "toolConfirmation": {
-                            "hint": "test hint for request_confirmation",
+                            "hint": HINT_TEXT,
                             "confirmed": False,
                         },
                     },
                 )
             ),
         ),
+        (
+            sub_agent1.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=sub_agent1.tools[0].name,
+                    response=TOOL_CALL_ERROR_RESPONSE,
+                )
+            ),
+        ),
     ]
     assert testing_utils.simplify_resumable_app_events(
         copy.deepcopy(sub_agent2_branch_events)
@@ -807,13 +835,22 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
                             "args": {},
                         },
                         "toolConfirmation": {
-                            "hint": "test hint for request_confirmation",
+                            "hint": HINT_TEXT,
                             "confirmed": False,
                         },
                     },
                 )
             ),
         ),
+        (
+            sub_agent2.name,
+            Part(
+                function_response=FunctionResponse(
+                    name=sub_agent2.tools[0].name,
+                    response=TOOL_CALL_ERROR_RESPONSE,
+                )
+            ),
+        ),
     ]
 
     ask_for_confirmation_function_call_ids = [
@@ -893,7 +930,7 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
             sub_agent2.name,
             Part(
                 function_response=FunctionResponse(
-                    name=sub_agent1.tools[0].name,
+                    name=sub_agent2.tools[0].name,
                     response={"result": "confirmed=True"},
                 )
             ),