feat: Implement checkpoint and resume logic for SequentialAgent

PiperOrigin-RevId: 811977004

Author: DeanChensj

src/google/adk/agents/base_agent.py

@@ -163,44 +163,41 @@ def _load_agent_state(
  self,
  ctx: InvocationContext,
  state_type: Type[AgentState],
- default_state: AgentState,
- ) -> tuple[AgentState, bool]:
+ ) -> Optional[AgentState]:
  """Loads the agent state from the invocation context, handling resumption.
 
  Args:
  ctx: The invocation context.
  state_type: The type of the agent state.
- default_state: The default state to use if not resuming.
 
  Returns:
- tuple[AgentState, bool]: The current state and a boolean indicating if
- resuming.
+ The current state if resuming, otherwise None.
  """
+ if not ctx.is_resumable:
+ return None
+
  if self.name not in ctx.agent_states:
- return default_state, False
+ return None
  else:
- return state_type.model_validate(ctx.agent_states.get(self.name)), True
+ return state_type.model_validate(ctx.agent_states.get(self.name))
 
  def _create_agent_state_event(
  self,
  ctx: InvocationContext,
  *,
- state: Optional[BaseAgentState] = None,
+ agent_state: Optional[BaseAgentState] = None,
  end_of_agent: bool = False,
  ) -> Event:
- """Creates an event for agent state.
+ """Returns an event with agent state.
 
  Args:
  ctx: The invocation context.
- state: The agent state to checkpoint.
+ agent_state: The agent state to checkpoint.
  end_of_agent: Whether the agent is finished running.
-
- Returns:
- An Event object representing the checkpoint.
  """
  event_actions = EventActions()
- if state:
- event_actions.agent_state = state.model_dump(mode='json')
+ if agent_state:
+ event_actions.agent_state = agent_state.model_dump(mode='json')
  if end_of_agent:
  event_actions.end_of_agent = True
  return Event(

src/google/adk/agents/invocation_context.py

@@ -208,6 +208,14 @@ class InvocationContext(BaseModel):
  of this invocation.
  """
 
+ @property
+ def is_resumable(self) -> bool:
+ """Returns whether the current invocation is resumable."""
+ return (
+ self.resumability_config is not None
+ and self.resumability_config.is_resumable
+ )
+
  def reset_agent_state(self, agent_name: str) -> None:
  """Resets the state of an agent, allowing it to be re-run."""
  self.agent_states.pop(agent_name, None)
@@ -284,10 +292,7 @@ def should_pause_invocation(self, event: Event) -> bool:
  Returns:
  Whether to pause the invocation right after this event.
  """
- if (
- not self.resumability_config
- or not self.resumability_config.is_resumable
- ):
+ if not self.is_resumable:
  return False
 
  if not event.long_running_tool_ids or not event.get_function_calls():

src/google/adk/agents/sequential_agent.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import logging
 from typing import AsyncGenerator
 from typing import ClassVar
 from typing import Type
@@ -32,6 +33,8 @@
 from .llm_agent import LlmAgent
 from .sequential_agent_config import SequentialAgentConfig
 
+logger = logging.getLogger('google_adk.' + __name__)
+
 
 @experimental
 class SequentialAgentState(BaseAgentState):
@@ -51,24 +54,69 @@ class SequentialAgent(BaseAgent):
  async def _run_async_impl(
  self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
- # Skip if there is no sub-agent.
  if not self.sub_agents:
  return
 
- for sub_agent in self.sub_agents:
- pause_invocation = False
+ # Initialize or resume the execution state from the agent state.
+ agent_state = self._load_agent_state(ctx, SequentialAgentState)
+ start_index = self._get_start_index(agent_state)
+
+ pause_invocation = False
+ resuming_sub_agent = agent_state is not None
+ for i in range(start_index, len(self.sub_agents)):
+ sub_agent = self.sub_agents[i]
+ if not resuming_sub_agent:
+ # If we are resuming from the current event, it means the same event has
+ # already been logged, so we should avoid yielding it again.
+ if ctx.is_resumable:
+ agent_state = SequentialAgentState(current_sub_agent=sub_agent.name)
+ yield self._create_agent_state_event(ctx, agent_state=agent_state)
+
+ # Reset the sub-agent's state in the context to ensure that each
+ # sub-agent starts fresh.
+ ctx.reset_agent_state(sub_agent.name)
 
  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.
+ # Skip the rest of the sub-agents if the invocation is paused.
  if pause_invocation:
  return
 
+ # Reset the flag for the next sub-agent.
+ resuming_sub_agent = False
+
+ if ctx.is_resumable:
+ yield self._create_agent_state_event(ctx, end_of_agent=True)
+
+ def _get_start_index(
+ self,
+ agent_state: SequentialAgentState,
+ ) -> int:
+ """Calculates the start index for the sub-agent loop."""
+ if not agent_state:
+ return 0
+
+ if not agent_state.current_sub_agent:
+ # This means the process was finished.
+ return len(self.sub_agents)
+
+ try:
+ sub_agent_names = [sub_agent.name for sub_agent in self.sub_agents]
+ return sub_agent_names.index(agent_state.current_sub_agent)
+ except ValueError:
+ # A sub-agent was removed so the agent name is not found.
+ # For now, we restart from the beginning.
+ logger.warning(
+ 'Sub-agent %s was removed so the agent name is not found. Restarting'
+ ' from the beginning.',
+ agent_state.current_sub_agent,
+ )
+ return 0
+
  @override
  async def _run_live_impl(
  self, ctx: InvocationContext

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

@@ -350,6 +350,8 @@ def _get_current_turn_contents(
  # Find the latest event that starts the current turn and process from there
  for i in range(len(events) - 1, -1, -1):
  event = events[i]
+ if not event.content:
+ continue
  if event.author == 'user' or _is_other_agent_reply(agent_name, event):
  return _get_contents(current_branch, events[i:], agent_name)
 

tests/unittests/agents/test_base_agent.py

@@ -26,6 +26,7 @@
 from google.adk.agents.base_agent import BaseAgentState
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.invocation_context import InvocationContext
+from google.adk.apps.app import ResumabilityConfig
 from google.adk.events.event import Event
 from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.plugins.plugin_manager import PluginManager
@@ -733,39 +734,6 @@ async def test_run_live_incomplete_agent(request: pytest.FixtureRequest):
  [e async for e in agent.run_live(parent_ctx)]
 
 
-@pytest.mark.asyncio
-async def test_create_agent_state_event(request: pytest.FixtureRequest):
- # Arrange
- agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
- ctx = await _create_parent_invocation_context(
- request.function.__name__, agent, branch='test_branch'
- )
- state = BaseAgentState()
-
- # Act
- event = agent._create_agent_state_event(ctx, state=state)
-
- # Assert
- assert event.invocation_id == ctx.invocation_id
- assert event.author == agent.name
- assert event.branch == 'test_branch'
- assert event.actions is not None
- assert event.actions.agent_state is not None
- assert event.actions.agent_state == state.model_dump(mode='json')
- assert not event.actions.end_of_agent
-
- # Act
- event = agent._create_agent_state_event(ctx, end_of_agent=True)
-
- # Assert
- assert event.invocation_id == ctx.invocation_id
- assert event.author == agent.name
- assert event.branch == 'test_branch'
- assert event.actions is not None
- assert event.actions.end_of_agent
- assert event.actions.agent_state is None
-
-
 def test_set_parent_agent_for_sub_agents(request: pytest.FixtureRequest):
  sub_agents: list[BaseAgent] = [
  _TestingAgent(name=f'{request.function.__name__}_sub_agent_1'),
@@ -895,7 +863,7 @@ class _TestAgentState(BaseAgentState):
 
 
 @pytest.mark.asyncio
-async def test_load_agent_state_no_resume():
+async def test_load_agent_state_not_resumable():
  agent = BaseAgent(name='test_agent')
  session_service = InMemorySessionService()
  session = await session_service.create_session(
@@ -907,14 +875,15 @@ async def test_load_agent_state_no_resume():
  session=session,
  session_service=session_service,
  )
- default_state = _TestAgentState(test_field='default')
 
- state, is_resuming = agent._load_agent_state(
-  ctx, _TestAgentState, default_state
- )
+ # Test case 1: resumability_config is None
+ state = agent._load_agent_state(ctx, _TestAgentState)
+ assert state is None
 
- assert not is_resuming
- assert state == default_state
+ # Test case 2: is_resumable is False
+ ctx.resumability_config = ResumabilityConfig(is_resumable=False)
+ state = agent._load_agent_state(ctx, _TestAgentState)
+ assert state is None
 
 
 @pytest.mark.asyncio
@@ -929,13 +898,70 @@ async def test_load_agent_state_with_resume():
  agent=agent,
  session=session,
  session_service=session_service,
+ resumability_config=ResumabilityConfig(is_resumable=True),
  )
+
+ # Test case 1: agent state not in context
+ state = agent._load_agent_state(ctx, _TestAgentState)
+ assert state is None
+
+ # Test case 2: agent state in context
  persisted_state = _TestAgentState(test_field='resumed')
  ctx.agent_states[agent.name] = persisted_state.model_dump(mode='json')
 
- state, is_resuming = agent._load_agent_state(
- ctx, _TestAgentState, _TestAgentState()
+ state = agent._load_agent_state(ctx, _TestAgentState)
+ assert state == persisted_state
+
+
+@pytest.mark.asyncio
+async def test_create_agent_state_event():
+ agent = BaseAgent(name='test_agent')
+ session_service = InMemorySessionService()
+ session = await session_service.create_session(
+ app_name='test_app', user_id='test_user'
+ )
+ ctx = InvocationContext(
+ invocation_id='test_invocation',
+ agent=agent,
+ session=session,
+ session_service=session_service,
  )
 
- assert is_resuming
- assert state == persisted_state
+ ctx.branch = 'test_branch'
+
+ # Test case 1: with state
+ state = _TestAgentState(test_field='checkpoint')
+ event = agent._create_agent_state_event(ctx, agent_state=state)
+ assert event is not None
+ assert event.invocation_id == ctx.invocation_id
+ assert event.author == agent.name
+ assert event.branch == 'test_branch'
+ assert event.actions is not None
+ assert event.actions.agent_state is not None
+ assert event.actions.agent_state == state.model_dump(mode='json')
+ assert not event.actions.end_of_agent
+
+ # Test case 2: with end_of_agent
+ event = agent._create_agent_state_event(ctx, end_of_agent=True)
+ assert event is not None
+ assert event.invocation_id == ctx.invocation_id
+ assert event.author == agent.name
+ assert event.branch == 'test_branch'
+ assert event.actions is not None
+ assert event.actions.end_of_agent
+ assert event.actions.agent_state is None
+
+ # Test case 3: with both state and end_of_agent
+ state = _TestAgentState(test_field='checkpoint')
+ event = agent._create_agent_state_event(
+ ctx, agent_state=state, end_of_agent=True
+ )
+ assert event is not None
+ assert event.actions.agent_state == state.model_dump(mode='json')
+ assert event.actions.end_of_agent
+
+ # Test case 4: with neither
+ event = agent._create_agent_state_event(ctx)
+ assert event is not None
+ assert event.actions.agent_state is None
+ assert not event.actions.end_of_agent

tests/unittests/agents/test_invocation_context.py

@@ -206,3 +206,22 @@ def test_should_not_pause_invocation_with_no_function_calls(
  assert not mock_invocation_context.should_pause_invocation(
  nonpausable_event
  )
+
+ def test_is_resumable_true(self):
+ """Tests that is_resumable is True when resumability is enabled."""
+ invocation_context = self._create_test_invocation_context(
+ ResumabilityConfig(is_resumable=True)
+ )
+ assert invocation_context.is_resumable
+
+ def test_is_resumable_false(self):
+ """Tests that is_resumable is False when resumability is disabled."""
+ invocation_context = self._create_test_invocation_context(
+ ResumabilityConfig(is_resumable=False)
+ )
+ assert not invocation_context.is_resumable
+
+ def test_is_resumable_no_config(self):
+ """Tests that is_resumable is False when no resumability config is set."""
+ invocation_context = self._create_test_invocation_context(None)
+ assert not invocation_context.is_resumable

tests/unittests/agents/test_llm_agent_include_contents.py

@@ -219,9 +219,10 @@ async def test_include_contents_none_sequential_agents():
  runner = testing_utils.InMemoryRunner(sequential_agent)
  events = runner.run("Original user request")
 
- assert len(events) == 2
- assert events[0].author == "agent1"
- assert events[1].author == "agent2"
+ simplified_events = [event for event in events if event.content]
+ assert len(simplified_events) == 2
+ assert simplified_events[0].author == "agent1"
+ assert simplified_events[1].author == "agent2"
 
  # Agent1 sees original user request
  agent1_contents = testing_utils.simplify_contents(