feat: Modify runner to support resuming an invocation (optionally with a function response)

PiperOrigin-RevId: 813008406

Author: XinranTang

src/google/adk/agents/invocation_context.py

@@ -34,6 +34,7 @@
 from ..sessions.session import Session
 from .active_streaming_tool import ActiveStreamingTool
 from .base_agent import BaseAgent
+from .base_agent import BaseAgentState
 from .context_cache_config import ContextCacheConfig
 from .live_request_queue import LiveRequestQueue
 from .run_config import RunConfig
@@ -221,6 +222,37 @@ def reset_agent_state(self, agent_name: str) -> None:
     self.agent_states.pop(agent_name, None)
     self.end_of_agents.pop(agent_name, None)
 
+  def populate_invocation_agent_states(self) -> None:
+    """Populates agent states for the current invocation if it is resumable.
+
+    For history events that contain agent state information, set the
+    agent_state and end_of_agent of the agent that generated the event.
+
+    For non-workflow agents, also set an initial agent_state if it has
+    already generated some contents.
+    """
+    if not self.is_resumable:
+      return
+    for event in self._get_events(current_invocation=True):
+      if event.actions.end_of_agent:
+        self.end_of_agents[event.author] = True
+        # Delete agent_state when it is end
+        self.agent_states.pop(event.author, None)
+      elif event.actions.agent_state is not None:
+        self.agent_states[event.author] = event.actions.agent_state
+        # Invalidate the end_of_agent flag
+        self.end_of_agents[event.author] = False
+      elif (
+          event.author != "user"
+          and event.content
+          and not self.agent_states.get(event.author)
+      ):
+        # If the agent has generated some contents but its agent_state is not
+        # set, set its agent_state to an empty agent_state.
+        self.agent_states[event.author] = BaseAgentState()
+        # Invalidate the end_of_agent flag
+        self.end_of_agents[event.author] = False
+
   def increment_llm_call_count(
       self,
   ):

src/google/adk/runners.py

@@ -29,6 +29,7 @@
 
 from .agents.active_streaming_tool import ActiveStreamingTool
 from .agents.base_agent import BaseAgent
+from .agents.base_agent import BaseAgentState
 from .agents.context_cache_config import ContextCacheConfig
 from .agents.invocation_context import InvocationContext
 from .agents.invocation_context import new_invocation_context_id
@@ -272,7 +273,8 @@ async def run_async(
       *,
       user_id: str,
       session_id: str,
-      new_message: types.Content,
+      invocation_id: Optional[str] = None,
+      new_message: Optional[types.Content] = None,
       state_delta: Optional[dict[str, Any]] = None,
       run_config: Optional[RunConfig] = None,
   ) -> AsyncGenerator[Event, None]:
@@ -281,6 +283,8 @@ async def run_async(
     Args:
       user_id: The user ID of the session.
       session_id: The session ID of the session.
+      invocation_id: The invocation ID of the session, set this to resume an
+        interrupted invocation.
       new_message: A new message to append to the session.
       state_delta: Optional state changes to apply to the session.
       run_config: The run config for the agent.
@@ -289,29 +293,57 @@ async def run_async(
       The events generated by the agent.
 
     Raises:
-      ValueError: If the session is not found.
+      ValueError: If the session is not found; If both invocation_id and
+        new_message are None.
     """
     run_config = run_config or RunConfig()
 
-    if not new_message.role:
+    if new_message and not new_message.role:
       new_message.role = 'user'
 
     async def _run_with_trace(
-        new_message: types.Content,
+        new_message: Optional[types.Content] = None,
+        invocation_id: Optional[str] = None,
     ) -> AsyncGenerator[Event, None]:
       with tracer.start_as_current_span('invocation'):
         session = await self.session_service.get_session(
             app_name=self.app_name, user_id=user_id, session_id=session_id
         )
         if not session:
           raise ValueError(f'Session not found: {session_id}')
-
-        invocation_context = await self._setup_context_for_new_invocation(
-            session=session,
-            new_message=new_message,
-            run_config=run_config,
-            state_delta=state_delta,
-        )
+        if not invocation_id and not new_message:
+          raise ValueError('Both invocation_id and new_message are None.')
+
+        if invocation_id:
+          if (
+              not self.resumability_config
+              or not self.resumability_config.is_resumable
+          ):
+            raise ValueError(
+                f'invocation_id: {invocation_id} is provided but the app is not'
+                ' resumable.'
+            )
+          invocation_context = await self._setup_context_for_resumed_invocation(
+              session=session,
+              new_message=new_message,
+              invocation_id=invocation_id,
+              run_config=run_config,
+              state_delta=state_delta,
+          )
+          if invocation_context.end_of_agents.get(self.agent.name):
+            # Directly return if the root agent has already ended.
+            # TODO: Handle the case where the invocation-to-resume started from
+            # a sub_agent:
+            # invocation1: root_agent -> sub_agent1
+            # invocation2: sub_agent1 [paused][resume]
+            return
+        else:
+          invocation_context = await self._setup_context_for_new_invocation(
+              session=session,
+              new_message=new_message,  # new_message is not None.
+              run_config=run_config,
+              state_delta=state_delta,
+          )
 
         async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
           async with Aclosing(ctx.agent.run_async(ctx)) as agen:
@@ -329,7 +361,7 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
           async for event in agen:
             yield event
 
-    async with Aclosing(_run_with_trace(new_message)) as agen:
+    async with Aclosing(_run_with_trace(new_message, invocation_id)) as agen:
       async for event in agen:
         yield event
 
@@ -462,6 +494,11 @@ async def _append_new_message_to_session(
           author='user',
           content=new_message,
       )
+    # If new_message is a function response, find the matching function call
+    # and use its branch as the new event's branch.
+    if function_call := invocation_context._find_matching_function_call(event):
+      event.branch = function_call.branch
+
     await self.session_service.append_event(session=session, event=event)
 
   async def run_live(
@@ -692,10 +729,82 @@ async def _setup_context_for_new_invocation(
     invocation_context.agent = self._find_agent_to_run(session, self.agent)
     return invocation_context
 
+  async def _setup_context_for_resumed_invocation(
+      self,
+      *,
+      session: Session,
+      new_message: Optional[types.Content],
+      invocation_id: Optional[str],
+      run_config: RunConfig,
+      state_delta: Optional[dict[str, Any]],
+  ) -> InvocationContext:
+    """Sets up the context for a resumed invocation.
+
+    Args:
+      session: The session to setup the invocation context for.
+      new_message: The new message to process and append to the session.
+      invocation_id: The invocation id to resume.
+      run_config: The run config of the agent.
+      state_delta: Optional state changes to apply to the session.
+
+    Returns:
+      The invocation context for the resumed invocation.
+
+    Raises:
+      ValueError: If the session has no events to resume; If no user message is
+        available for resuming the invocation; Or if the app is not resumable.
+    """
+    if not session.events:
+      raise ValueError(f'Session {session.id} has no events to resume.')
+
+    # Step 1: Maybe retrive a previous user message for the invocation.
+    user_message = new_message or self._find_user_message_for_invocation(
+        session.events, invocation_id
+    )
+    if not user_message:
+      raise ValueError(
+          f'No user message available for resuming invocation: {invocation_id}'
+      )
+    # Step 2: Create invocation context.
+    invocation_context = self._new_invocation_context(
+        session,
+        new_message=user_message,
+        run_config=run_config,
+        invocation_id=invocation_id,
+    )
+    # Step 3: Maybe handle new message.
+    if new_message:
+      await self._handle_new_message(
+          session=session,
+          new_message=user_message,
+          invocation_context=invocation_context,
+          run_config=run_config,
+          state_delta=state_delta,
+      )
+    # Step 4: Populate agent states for the current invocation.
+    invocation_context.populate_invocation_agent_states()
+    return invocation_context
+
+  def _find_user_message_for_invocation(
+      self, events: list[Event], invocation_id: str
+  ) -> Optional[types.Content]:
+    """Finds the user message that started a specific invocation."""
+    for event in events:
+      if (
+          event.invocation_id == invocation_id
+          and event.author == 'user'
+          and event.content
+          and event.content.parts
+          and event.content.parts[0].text
+      ):
+        return event.content
+    return None
+
   def _new_invocation_context(
       self,
       session: Session,
       *,
+      invocation_id: Optional[str] = None,
       new_message: Optional[types.Content] = None,
       live_request_queue: Optional[LiveRequestQueue] = None,
       run_config: Optional[RunConfig] = None,
@@ -704,6 +813,7 @@ def _new_invocation_context(
 
     Args:
         session: The session for the context.
+        invocation_id: The invocation id for the context.
         new_message: The new message for the context.
         live_request_queue: The live request queue for the context.
         run_config: The run config for the context.
@@ -712,7 +822,7 @@ def _new_invocation_context(
         The new invocation context.
     """
     run_config = run_config or RunConfig()
-    invocation_id = new_invocation_context_id()
+    invocation_id = invocation_id or new_invocation_context_id()
 
     if run_config.support_cfc and isinstance(self.agent, LlmAgent):
       model_name = self.agent.canonical_model.model

tests/unittests/agents/test_invocation_context.py

@@ -15,9 +15,11 @@
 from unittest.mock import Mock
 
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.base_agent import BaseAgentState
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.apps import ResumabilityConfig
 from google.adk.events.event import Event
+from google.adk.events.event_actions import EventActions
 from google.adk.sessions.base_session_service import BaseSessionService
 from google.adk.sessions.session import Session
 from google.genai.types import Content
@@ -227,6 +229,124 @@ def test_is_resumable_no_config(self):
     invocation_context = self._create_test_invocation_context(None)
     assert not invocation_context.is_resumable
 
+  def test_populate_invocation_agent_states_not_resumable(self):
+    """Tests that populate_invocation_agent_states does nothing if not resumable."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=False)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=True, agent_state=None),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert not invocation_context.end_of_agents
+
+  def test_populate_invocation_agent_states_end_of_agent(self):
+    """Tests that populate_invocation_agent_states handles end_of_agent."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=True, agent_state=None),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert invocation_context.end_of_agents == {'agent1': True}
+
+  def test_populate_invocation_agent_states_with_agent_state(self):
+    """Tests that populate_invocation_agent_states handles agent_state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(
+            end_of_agent=False,
+            agent_state=BaseAgentState().model_dump(mode='json'),
+        ),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert invocation_context.agent_states == {'agent1': {}}
+    assert invocation_context.end_of_agents == {'agent1': False}
+
+  def test_populate_invocation_agent_states_with_agent_state_and_end_of_agent(
+      self,
+  ):
+    """Tests that populate_invocation_agent_states handles agent_state and end_of_agent."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(
+            end_of_agent=True,
+            agent_state=BaseAgentState().model_dump(mode='json'),
+        ),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    # When both agent_state and end_of_agent are set, agent_state should be
+    # cleared, as end_of_agent is of a higher priority.
+    assert not invocation_context.agent_states
+    assert invocation_context.end_of_agents == {'agent1': True}
+
+  def test_populate_invocation_agent_states_with_content_no_state(self):
+    """Tests that populate_invocation_agent_states creates default state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=False, agent_state=None),
+        content=Content(role='model', parts=[Part(text='hi')]),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert invocation_context.agent_states == {'agent1': BaseAgentState()}
+    assert invocation_context.end_of_agents == {'agent1': False}
+
+  def test_populate_invocation_agent_states_user_message_event(self):
+    """Tests that populate_invocation_agent_states ignores user message events for default state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='user',
+        actions=EventActions(end_of_agent=False, agent_state=None),
+        content=Content(role='user', parts=[Part(text='hi')]),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert not invocation_context.end_of_agents
+
+  def test_populate_invocation_agent_states_no_content(self):
+    """Tests that populate_invocation_agent_states ignores events with no content if no state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=None, agent_state=None),
+        content=None,
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert not invocation_context.end_of_agents
+
 
 class TestFindMatchingFunctionCall:
   """Test suite for find_matching_function_call."""