feat: Implement checkpoint and resume logic for ParallelAgent

PiperOrigin-RevId: 812658378

Author: DeanChensj

src/google/adk/agents/parallel_agent.py

@@ -26,6 +26,7 @@
 from ..events.event import Event
 from ..utils.context_utils import Aclosing
 from .base_agent import BaseAgent
+from .base_agent import BaseAgentState
 from .base_agent_config import BaseAgentConfig
 from .invocation_context import InvocationContext
 from .parallel_agent_config import ParallelAgentConfig
@@ -178,12 +179,22 @@ async def _run_async_impl(
     if not self.sub_agents:
       return
 
-    agent_runs = [
-        sub_agent.run_async(
-            _create_branch_ctx_for_sub_agent(self, sub_agent, ctx)
-        )
-        for sub_agent in self.sub_agents
-    ]
+    agent_state = self._load_agent_state(ctx, BaseAgentState)
+    if ctx.is_resumable and agent_state is None:
+      yield self._create_agent_state_event(ctx, agent_state=BaseAgentState())
+
+    agent_runs = []
+    # Prepare and collect async generators for each sub-agent.
+    for sub_agent in self.sub_agents:
+      if agent_state is None:
+        # Reset sub-agent state to make sure each sub-agent starts fresh.
+        ctx.reset_agent_state(sub_agent.name)
+
+      sub_agent_ctx = _create_branch_ctx_for_sub_agent(self, sub_agent, ctx)
+
+      # Only include sub-agents that haven't finished in a previous run.
+      if not sub_agent_ctx.end_of_agents.get(sub_agent.name):
+        agent_runs.append(sub_agent.run_async(sub_agent_ctx))
 
     pause_invocation = False
     try:
@@ -203,6 +214,10 @@ async def _run_async_impl(
       if pause_invocation:
         return
 
+      # Once all sub-agents are done, mark the ParallelAgent as final.
+      if ctx.is_resumable:
+        yield self._create_agent_state_event(ctx, end_of_agent=True)
+
     finally:
       for sub_agent_run in agent_runs:
         await sub_agent_run.aclose()

tests/unittests/agents/test_parallel_agent.py

@@ -18,9 +18,12 @@
 from typing import AsyncGenerator
 
 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.agents.parallel_agent import ParallelAgent
 from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.agents.sequential_agent import SequentialAgentState
+from google.adk.apps.app import ResumabilityConfig
 from google.adk.events.event import Event
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.genai import types
@@ -52,7 +55,7 @@ async def _run_async_impl(
 
 
 async def _create_parent_invocation_context(
-    test_name: str, agent: BaseAgent
+    test_name: str, agent: BaseAgent, is_resumable: bool = False
 ) -> InvocationContext:
   session_service = InMemorySessionService()
   session = await session_service.create_session(
@@ -63,11 +66,13 @@ async def _create_parent_invocation_context(
       agent=agent,
       session=session,
       session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=is_resumable),
   )
 
 
 @pytest.mark.asyncio
-async def test_run_async(request: pytest.FixtureRequest):
+@pytest.mark.parametrize('is_resumable', [True, False])
+async def test_run_async(request: pytest.FixtureRequest, is_resumable: bool):
   agent1 = _TestingAgent(
       name=f'{request.function.__name__}_test_agent_1',
       delay=0.5,
@@ -81,23 +86,43 @@ async def test_run_async(request: pytest.FixtureRequest):
       ],
   )
   parent_ctx = await _create_parent_invocation_context(
-      request.function.__name__, parallel_agent
+      request.function.__name__, parallel_agent, is_resumable=is_resumable
   )
   events = [e async for e in parallel_agent.run_async(parent_ctx)]
 
-  assert len(events) == 2
-  # agent2 generates an event first, then agent1. Because they run in parallel
-  # and agent1 has a delay.
-  assert events[0].author == agent2.name
-  assert events[1].author == agent1.name
-  assert events[0].branch.endswith(f'{parallel_agent.name}.{agent2.name}')
-  assert events[1].branch.endswith(f'{parallel_agent.name}.{agent1.name}')
-  assert events[0].content.parts[0].text == f'Hello, async {agent2.name}!'
-  assert events[1].content.parts[0].text == f'Hello, async {agent1.name}!'
+  if is_resumable:
+    assert len(events) == 4
+
+    assert events[0].author == parallel_agent.name
+    assert not events[0].actions.end_of_agent
+
+    # agent2 generates an event first, then agent1. Because they run in parallel
+    # and agent1 has a delay.
+    assert events[1].author == agent2.name
+    assert events[2].author == agent1.name
+    assert events[1].branch == f'{parallel_agent.name}.{agent2.name}'
+    assert events[2].branch == f'{parallel_agent.name}.{agent1.name}'
+    assert events[1].content.parts[0].text == f'Hello, async {agent2.name}!'
+    assert events[2].content.parts[0].text == f'Hello, async {agent1.name}!'
+
+    assert events[3].author == parallel_agent.name
+    assert events[3].actions.end_of_agent
+  else:
+    assert len(events) == 2
+
+    assert events[0].author == agent2.name
+    assert events[1].author == agent1.name
+    assert events[0].branch == f'{parallel_agent.name}.{agent2.name}'
+    assert events[1].branch == f'{parallel_agent.name}.{agent1.name}'
+    assert events[0].content.parts[0].text == f'Hello, async {agent2.name}!'
+    assert events[1].content.parts[0].text == f'Hello, async {agent1.name}!'
 
 
 @pytest.mark.asyncio
-async def test_run_async_branches(request: pytest.FixtureRequest):
+@pytest.mark.parametrize('is_resumable', [True, False])
+async def test_run_async_branches(
+    request: pytest.FixtureRequest, is_resumable: bool
+):
   agent1 = _TestingAgent(
       name=f'{request.function.__name__}_test_agent_1',
       delay=0.5,
@@ -116,28 +141,124 @@ async def test_run_async_branches(request: pytest.FixtureRequest):
       ],
   )
   parent_ctx = await _create_parent_invocation_context(
-      request.function.__name__, parallel_agent
+      request.function.__name__, parallel_agent, is_resumable=is_resumable
   )
   events = [e async for e in parallel_agent.run_async(parent_ctx)]
 
-  assert len(events) == 3
-  assert (
-      events[0].author == agent2.name
-      and events[0].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+  if is_resumable:
+    assert len(events) == 8
+
+    # 1. parallel agent checkpoint
+    assert events[0].author == parallel_agent.name
+    assert not events[0].actions.end_of_agent
+
+    # 2. sequential agent checkpoint
+    assert events[1].author == sequential_agent.name
+    assert not events[1].actions.end_of_agent
+    assert events[1].actions.agent_state['current_sub_agent'] == agent2.name
+    assert events[1].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # 3. agent 2 event
+    assert events[2].author == agent2.name
+    assert events[2].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # 4. sequential agent checkpoint
+    assert events[3].author == sequential_agent.name
+    assert not events[3].actions.end_of_agent
+    assert events[3].actions.agent_state['current_sub_agent'] == agent3.name
+    assert events[3].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # 5. agent 3 event
+    assert events[4].author == agent3.name
+    assert events[4].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # 6. sequential agent checkpoint (end)
+    assert events[5].author == sequential_agent.name
+    assert events[5].actions.end_of_agent
+    assert events[5].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # Descendants of the same sub-agent should have the same branch.
+    assert events[1].branch == events[2].branch
+    assert events[2].branch == events[3].branch
+    assert events[3].branch == events[4].branch
+    assert events[4].branch == events[5].branch
+
+    # 7. agent 1 event
+    assert events[6].author == agent1.name
+    assert events[6].branch == f'{parallel_agent.name}.{agent1.name}'
+
+    # Sub-agents should have different branches.
+    assert events[6].branch != events[1].branch
+
+    # 8. parallel agent checkpoint (end)
+    assert events[7].author == parallel_agent.name
+    assert events[7].actions.end_of_agent
+  else:
+    assert len(events) == 3
+
+    # 1. agent 2 event
+    assert events[0].author == agent2.name
+    assert events[0].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # 2. agent 3 event
+    assert events[1].author == agent3.name
+    assert events[1].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+    # 3. agent 1 event
+    assert events[2].author == agent1.name
+    assert events[2].branch == f'{parallel_agent.name}.{agent1.name}'
+
+
+@pytest.mark.asyncio
+async def test_resume_async_branches(request: pytest.FixtureRequest):
+  agent1 = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent_1', delay=0.5
+  )
+  agent2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  agent3 = _TestingAgent(name=f'{request.function.__name__}_test_agent_3')
+  sequential_agent = SequentialAgent(
+      name=f'{request.function.__name__}_test_sequential_agent',
+      sub_agents=[agent2, agent3],
+  )
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[
+          sequential_agent,
+          agent1,
+      ],
   )
-  assert (
-      events[1].author == agent3.name
-      and events[0].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent, is_resumable=True
   )
-  # Descendants of the same sub-agent should have the same branch.
-  assert events[0].branch == events[1].branch
-  assert (
-      events[2].author == agent1.name
-      and events[2].branch == f'{parallel_agent.name}.{agent1.name}'
+  parent_ctx.agent_states[parallel_agent.name] = BaseAgentState().model_dump(
+      mode='json'
   )
-  # Sub-agents should have different branches.
-  assert events[2].branch != events[1].branch
-  assert events[2].branch != events[0].branch
+  parent_ctx.agent_states[sequential_agent.name] = SequentialAgentState(
+      current_sub_agent=agent3.name
+  ).model_dump(mode='json')
+
+  events = [e async for e in parallel_agent.run_async(parent_ctx)]
+
+  assert len(events) == 4
+
+  # The sequential agent resumes from agent3.
+  # 1. Agent 3 event
+  assert events[0].author == agent3.name
+  assert events[0].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+  # 2. Sequential agent checkpoint (end)
+  assert events[1].author == sequential_agent.name
+  assert events[1].actions.end_of_agent
+  assert events[1].branch == f'{parallel_agent.name}.{sequential_agent.name}'
+
+  # Agent 1 runs in parallel but has a delay.
+  # 3. Agent 1 event
+  assert events[2].author == agent1.name
+  assert events[2].branch == f'{parallel_agent.name}.{agent1.name}'
+
+  # 4. Parallel agent checkpoint (end)
+  assert events[3].author == parallel_agent.name
+  assert events[3].actions.end_of_agent
 
 
 class _TestingAgentWithMultipleEvents(_TestingAgent):