fix: Yield the long running tool response before pausing execution

PiperOrigin-RevId: 825056377

Author: DeanChensj

src/google/adk/agents/llm_agent.py

@@ -430,16 +430,24 @@ async def _run_async_impl(
       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 @@ async def _run_one_step_async(
     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 @@ async def run_async(
                 ' 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 @@ def test_pause_on_long_running_function_call(
     """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 @@ def test_pause_first_agent_on_long_running_function_call(
             ),
         ),
         ("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 @@ def test_pause_second_agent_on_long_running_function_call(
         (
             "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 @@ def test_pause_second_agent_on_long_running_function_call(
             ),
         ),
         ("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 @@ def test_pause_on_long_running_function_call(
             ),
         ),
         ("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 @@ def test_pause_on_long_running_function_call(
         ("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 @@ def test_pause_on_long_running_function_call(
         ("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"}
+            ),
+        ),
     ]