feat: Pause invocations on long running function calls for resumable apps

PiperOrigin-RevId: 811518771

Author: XinranTang

src/google/adk/agents/invocation_context.py

@@ -260,6 +260,45 @@ def get_events(
       results = [event for event in results if event.branch == self.branch]
     return results
 
+  def should_pause_invocation(self, event: Event) -> bool:
+    """Returns whether to pause the invocation right after this event.
+
+    "Pausing" an invocation is different from "ending" an invocation. A paused
+    invocation can be resumed later, while an ended invocation cannot.
+
+    Pausing the current agent's run will also pause all the agents that
+    depend on its execution, i.e. the subsequent agents in a workflow, and the
+    current agent's ancestors, etc.
+
+    Note that parallel sibling agents won't be affected, but their common
+    ancestors will be paused after all the non-blocking sub-agents finished
+    running.
+
+    Should meet all following conditions to pause an invocation:
+      1. The app is resumable.
+      2. The current event has a long running function call.
+
+    Args:
+      event: The current event.
+
+    Returns:
+      Whether to pause the invocation right after this event.
+    """
+    if (
+        not self.resumability_config
+        or not self.resumability_config.is_resumable
+    ):
+      return False
+
+    if not event.long_running_tool_ids or not event.get_function_calls():
+      return False
+
+    for fc in event.get_function_calls():
+      if fc.id in event.long_running_tool_ids:
+        return True
+
+    return False
+
 
 def new_invocation_context_id() -> str:
   return "e-" + str(uuid.uuid4())

src/google/adk/agents/llm_agent.py

@@ -341,6 +341,8 @@ async def _run_async_impl(
       async for event in agen:
         self.__maybe_save_output_to_state(event)
         yield event
+        if ctx.should_pause_invocation(event):
+          return
 
   @override
   async def _run_live_impl(

src/google/adk/agents/loop_agent.py

@@ -70,15 +70,26 @@ async def _run_async_impl(
     while not self.max_iterations or times_looped < self.max_iterations:
       for sub_agent in self.sub_agents:
         should_exit = False
+        pause_invocation = False
+
         async with Aclosing(sub_agent.run_async(ctx)) as agen:
           async for event in agen:
             yield event
             if event.actions.escalate:
               should_exit = True
+            if ctx.should_pause_invocation(event):
+              pause_invocation = True
 
+        # Indicates that the loop agent should exist after running this
+        # sub-agent.
         if should_exit:
           return
 
+        # Indicates that the invocation should be paused after running this
+        # sub-agent.
+        if pause_invocation:
+          return
+
       times_looped += 1
     return
 

src/google/adk/agents/parallel_agent.py

@@ -181,16 +181,26 @@ async def _run_async_impl(
         )
         for sub_agent in self.sub_agents
     ]
+
+    pause_invocation = False
     try:
       # TODO remove if once Python <3.11 is no longer supported.
       if sys.version_info >= (3, 11):
         async with Aclosing(_merge_agent_run(agent_runs)) as agen:
           async for event in agen:
             yield event
+            if ctx.should_pause_invocation(event):
+              pause_invocation = True
       else:
         async with Aclosing(_merge_agent_run_pre_3_11(agent_runs)) as agen:
           async for event in agen:
             yield event
+            if ctx.should_pause_invocation(event):
+              pause_invocation = True
+
+      if pause_invocation:
+        return
+
     finally:
       for sub_agent_run in agent_runs:
         await sub_agent_run.aclose()

src/google/adk/agents/sequential_agent.py

@@ -52,9 +52,18 @@ async def _run_async_impl(
       self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
     for sub_agent in self.sub_agents:
+      pause_invocation = False
+
       async with Aclosing(sub_agent.run_async(ctx)) as agen:
         async for event in agen:
           yield event
+          if ctx.should_pause_invocation(event):
+            pause_invocation = True
+
+      # Indicates the invocation should pause when receiving signal from
+      # the current sub_agent.
+      if pause_invocation:
+        return
 
   @override
   async def _run_live_impl(

tests/unittests/agents/test_invocation_context.py

@@ -16,11 +16,16 @@
 
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.invocation_context import InvocationContext
+from google.adk.apps import ResumabilityConfig
 from google.adk.events.event import Event
 from google.adk.sessions.base_session_service import BaseSessionService
 from google.adk.sessions.session import Session
+from google.genai.types import FunctionCall
+from google.genai.types import Part
 import pytest
 
+from .. import testing_utils
+
 
 class TestInvocationContext:
   """Test suite for InvocationContext."""
@@ -117,3 +122,87 @@ def test_get_events_with_no_matching_events(self, mock_invocation_context):
         current_branch=True,
     )
     assert not events
+
+
+class TestInvocationContextWithAppResumablity:
+  """Test suite for InvocationContext regarding app resumability."""
+
+  @pytest.fixture
+  def long_running_function_call(self) -> FunctionCall:
+    """A long running function call."""
+    return FunctionCall(
+        id='tool_call_id_1',
+        name='long_running_function_call',
+        args={},
+    )
+
+  @pytest.fixture
+  def event_to_pause(self, long_running_function_call) -> Event:
+    """An event with a long running function call."""
+    return Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=testing_utils.ModelContent(
+            [Part(function_call=long_running_function_call)]
+        ),
+        long_running_tool_ids=[long_running_function_call.id],
+    )
+
+  def _create_test_invocation_context(
+      self, resumability_config
+  ) -> InvocationContext:
+    """Create a mock invocation context for testing."""
+    ctx = InvocationContext(
+        session_service=Mock(spec=BaseSessionService),
+        agent=Mock(spec=BaseAgent),
+        invocation_id='inv_1',
+        session=Mock(spec=Session),
+        resumability_config=resumability_config,
+    )
+    return ctx
+
+  def test_should_pause_invocation_with_resumable_app(self, event_to_pause):
+    """Tests should_pause_invocation with a resumable app."""
+    mock_invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+
+    assert mock_invocation_context.should_pause_invocation(event_to_pause)
+
+  def test_should_not_pause_invocation_with_non_resumable_app(
+      self, event_to_pause
+  ):
+    """Tests should_pause_invocation with a non-resumable app."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=False)
+    )
+
+    assert not invocation_context.should_pause_invocation(event_to_pause)
+
+  def test_should_not_pause_invocation_with_no_long_running_tool_ids(
+      self, event_to_pause
+  ):
+    """Tests should_pause_invocation with no long running tools."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    nonpausable_event = event_to_pause.model_copy(
+        update={'long_running_tool_ids': []}
+    )
+
+    assert not invocation_context.should_pause_invocation(nonpausable_event)
+
+  def test_should_not_pause_invocation_with_no_function_calls(
+      self, event_to_pause
+  ):
+    """Tests should_pause_invocation with a non-model event."""
+    mock_invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    nonpausable_event = event_to_pause.model_copy(
+        update={'content': testing_utils.UserContent('test text part')}
+    )
+
+    assert not mock_invocation_context.should_pause_invocation(
+        nonpausable_event
+    )