feat: Make resumable llm agents yield checkpoint events

PiperOrigin-RevId: 813001108

src/google/adk/agents/invocation_context.py

@@ -242,8 +242,10 @@ class InvocationContext(BaseModel):
   def user_id(self) -> str:
     return self.session.user_id
 
-  def get_events(
+  # TODO: Move this method from invocation_context to a dedicated module.
+  def _get_events(
       self,
+      *,
       current_invocation: bool = False,
       current_branch: bool = False,
   ) -> list[Event]:
@@ -304,6 +306,25 @@ class InvocationContext(BaseModel):
 
     return False
 
+  # TODO: Move this method from invocation_context to a dedicated module.
+  # TODO: Converge this method with find_matching_function_call in llm_flows.
+  def _find_matching_function_call(
+      self, function_response_event: Event
+  ) -> Optional[Event]:
+    """Finds the function call event in the current invocation that matches the function response id."""
+    function_responses = function_response_event.get_function_responses()
+    if not function_responses:
+      return None
+    function_call_id = function_responses[0].id
+
+    events = self._get_events(current_invocation=True)
+    # The last event is function_response_event, so we search backwards from the
+    # one before it.
+    for event in reversed(events[:-1]):
+      if any(fc.id == function_call_id for fc in event.get_function_calls()):
+        return event
+    return None
+
 
 def new_invocation_context_id() -> str:
   return "e-" + str(uuid.uuid4())

src/google/adk/agents/llm_agent.py

@@ -54,6 +54,7 @@ from ..tools.tool_context import ToolContext
 from ..utils.context_utils import Aclosing
 from ..utils.feature_decorator import experimental
 from .base_agent import BaseAgent
+from .base_agent import BaseAgentState
 from .base_agent_config import BaseAgentConfig
 from .callback_context import CallbackContext
 from .invocation_context import InvocationContext
@@ -337,6 +338,20 @@ class LlmAgent(BaseAgent):
   async def _run_async_impl(
       self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
+    agent_state = self._load_agent_state(ctx, BaseAgentState)
+
+    # If there is an sub-agent to resume, run it and then end the current
+    # agent.
+    if agent_state is not None and (
+        agent_to_transfer := self._get_subagent_to_resume(ctx)
+    ):
+      async with Aclosing(agent_to_transfer.run_async(ctx)) as agen:
+        async for event in agen:
+          yield event
+
+      yield self._create_agent_state_event(ctx, end_of_agent=True)
+      return
+
     async with Aclosing(self._llm_flow.run_async(ctx)) as agen:
       async for event in agen:
         self.__maybe_save_output_to_state(event)
@@ -344,6 +359,9 @@ class LlmAgent(BaseAgent):
         if ctx.should_pause_invocation(event):
           return
 
+    if ctx.is_resumable:
+      yield self._create_agent_state_event(ctx, end_of_agent=True)
+
   @override
   async def _run_live_impl(
       self, ctx: InvocationContext
@@ -498,6 +516,74 @@ class LlmAgent(BaseAgent):
     else:
       return AutoFlow()
 
+  def _get_subagent_to_resume(
+      self, ctx: InvocationContext
+  ) -> Optional[BaseAgent]:
+    """Returns the sub-agent in the llm tree to resume if it exists.
+
+    There are 2 cases where we need to transfer to and resume a sub-agent:
+    1. The last event is a transfer to agent response from the current agent.
+       In this case, we need to return the agent specified in the response.
+
+    2. The last event's author isn't the current agent, or the user is
+       responding to another agent's tool call.
+       In this case, we need to return the LAST agent being transferred to
+       from the current agent.
+    """
+    events = ctx._get_events(current_invocation=True, current_branch=True)
+    if not events:
+      return None
+
+    last_event = events[-1]
+    if last_event.author == self.name:
+      # Last event is from current agent. Return transfer_to_agent in the event
+      # if it exists, or None.
+      return self.__get_transfer_to_agent_or_none(last_event, self.name)
+
+    # Last event is from user or another agent.
+    if last_event.author == 'user':
+      function_call_event = ctx._find_matching_function_call(last_event)
+      if not function_call_event:
+        raise ValueError(
+            'No agent to transfer to for resuming agent from function response'
+            f' {self.name}'
+        )
+      if function_call_event.author == self.name:
+        # User is responding to a tool call from the current agent.
+        # Current agent should continue, so no sub-agent to resume.
+        return None
+
+    # Last event is from another agent, or from user for another agent's tool
+    # call. We need to find the last agent we transferred to.
+    for event in reversed(events):
+      if agent := self.__get_transfer_to_agent_or_none(event, self.name):
+        return agent
+
+    return None
+
+  def __get_agent_to_run(self, agent_name: str) -> BaseAgent:
+    """Find the agent to run under the root agent by name."""
+    agent_to_run = self.root_agent.find_agent(agent_name)
+    if not agent_to_run:
+      raise ValueError(f'Agent {agent_name} not found in the agent tree.')
+    return agent_to_run
+
+  def __get_transfer_to_agent_or_none(
+      self, event: Event, from_agent: str
+  ) -> Optional[BaseAgent]:
+    """Returns the agent to run if the event is a transfer to agent response."""
+    function_responses = event.get_function_responses()
+    if not function_responses:
+      return None
+    for function_response in function_responses:
+      if (
+          function_response.name == 'transfer_to_agent'
+          and event.author == from_agent
+          and event.actions.transfer_to_agent != from_agent
+      ):
+        return self.__get_agent_to_run(event.actions.transfer_to_agent)
+    return None
+
   def __maybe_save_output_to_state(self, event: Event):
     """Saves the model output to state if needed."""
     # skip if the event was authored by some other agent (e.g. current agent

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

@@ -376,6 +376,28 @@ class BaseLlmFlow(ABC):
     if invocation_context.end_invocation:
       return
 
+    # Resume the LLM agent based on the last event from the current branch.
+    # 1. User content: continue the normal flow
+    # 2. Function call: call the tool and get the response event.
+    events = invocation_context._get_events(
+        current_invocation=True, current_branch=True
+    )
+    if (
+        invocation_context.is_resumable
+        and events
+        and events[-1].get_function_calls()
+    ):
+      model_response_event = events[-1]
+      async with Aclosing(
+          self._postprocess_handle_function_calls_async(
+              invocation_context, model_response_event, llm_request
+          )
+      ) as agen:
+        async for event in agen:
+          event.id = Event.new_id()
+          yield event
+        return
+
     # Calls the LLM.
     model_response_event = Event(
         id=Event.new_id(),

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

@@ -135,7 +135,8 @@ def _rearrange_events_for_latest_function_response(
   Returns:
     A list of events with the latest function_response rearranged.
   """
-  if not events:
+  if len(events) < 2:
+    # No need to process, since there is no function_call.
     return events
 
   function_responses = events[-1].get_function_responses()

src/google/adk/runners.py

@@ -606,7 +606,16 @@ class Runner:
     event = find_matching_function_call(session.events)
     if event and event.author:
       return root_agent.find_agent(event.author)
-    for event in filter(lambda e: e.author != 'user', reversed(session.events)):
+
+    def _event_filter(event: Event) -> bool:
+      """Filters out user-authored events and agent state change events."""
+      if event.author == 'user':
+        return False
+      if event.actions.agent_state is not None or event.actions.end_of_agent:
+        return False
+      return True
+
+    for event in filter(_event_filter, reversed(session.events)):
       if event.author == root_agent.name:
         # Found root agent.
         return root_agent

tests/unittests/agents/test_invocation_context.py

@@ -20,6 +20,7 @@ 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 Content
 from google.genai.types import FunctionCall
 from google.genai.types import Part
 import pytest
@@ -67,7 +68,7 @@ class TestInvocationContext:
       self, mock_invocation_context, mock_events
   ):
     """Tests that get_events returns all events when no filters are applied."""
-    events = mock_invocation_context.get_events()
+    events = mock_invocation_context._get_events()
     assert events == mock_events
 
   def test_get_events_filters_by_current_invocation(
@@ -75,7 +76,7 @@ class TestInvocationContext:
   ):
     """Tests that get_events correctly filters by the current invocation."""
     event1, event2, _, _ = mock_events
-    events = mock_invocation_context.get_events(current_invocation=True)
+    events = mock_invocation_context._get_events(current_invocation=True)
     assert events == [event1, event2]
 
   def test_get_events_filters_by_current_branch(
@@ -83,7 +84,7 @@ class TestInvocationContext:
   ):
     """Tests that get_events correctly filters by the current branch."""
     event1, _, event3, _ = mock_events
-    events = mock_invocation_context.get_events(current_branch=True)
+    events = mock_invocation_context._get_events(current_branch=True)
     assert events == [event1, event3]
 
   def test_get_events_filters_by_invocation_and_branch(
@@ -91,7 +92,7 @@ class TestInvocationContext:
   ):
     """Tests that get_events filters by invocation and branch."""
     event1, _, _, _ = mock_events
-    events = mock_invocation_context.get_events(
+    events = mock_invocation_context._get_events(
         current_invocation=True,
         current_branch=True,
     )
@@ -100,7 +101,7 @@ class TestInvocationContext:
   def test_get_events_with_no_events_in_session(self, mock_invocation_context):
     """Tests get_events when the session has no events."""
     mock_invocation_context.session.events = []
-    events = mock_invocation_context.get_events()
+    events = mock_invocation_context._get_events()
     assert not events
 
   def test_get_events_with_no_matching_events(self, mock_invocation_context):
@@ -109,15 +110,15 @@ class TestInvocationContext:
     mock_invocation_context.branch = 'branch_C'
 
     # Filter by invocation
-    events = mock_invocation_context.get_events(current_invocation=True)
+    events = mock_invocation_context._get_events(current_invocation=True)
     assert not events
 
     # Filter by branch
-    events = mock_invocation_context.get_events(current_branch=True)
+    events = mock_invocation_context._get_events(current_branch=True)
     assert not events
 
     # Filter by both
-    events = mock_invocation_context.get_events(
+    events = mock_invocation_context._get_events(
         current_invocation=True,
         current_branch=True,
     )
@@ -225,3 +226,114 @@ class TestInvocationContextWithAppResumablity:
     """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
+
+
+class TestFindMatchingFunctionCall:
+  """Test suite for find_matching_function_call."""
+
+  @pytest.fixture
+  def test_invocation_context(self):
+    """Create a mock invocation context for testing."""
+
+    def _create_invocation_context(events):
+      return InvocationContext(
+          session_service=Mock(spec=BaseSessionService),
+          agent=Mock(spec=BaseAgent, name='agent'),
+          invocation_id='inv_1',
+          session=Mock(spec=Session, events=events),
+      )
+
+    return _create_invocation_context
+
+  def test_find_matching_function_call_found(self, test_invocation_context):
+    """Tests that a matching function call is found."""
+    fc = Part.from_function_call(name='some_tool', args={})
+    fc.function_call.id = 'test_function_call_id'
+    fc_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=testing_utils.ModelContent([fc]),
+    )
+    fr = Part.from_function_response(
+        name='some_tool', response={'result': 'ok'}
+    )
+    fr.function_response.id = 'test_function_call_id'
+    fr_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=Content(role='user', parts=[fr]),
+    )
+    invocation_context = test_invocation_context([fc_event, fr_event])
+    matching_fc_event = invocation_context._find_matching_function_call(
+        fr_event
+    )
+    assert testing_utils.simplify_content(
+        matching_fc_event.content
+    ) == testing_utils.simplify_content(fc_event.content)
+
+  def test_find_matching_function_call_not_found(self, test_invocation_context):
+    """Tests that no matching function call is returned if id doesn't match."""
+    fc = Part.from_function_call(name='some_tool', args={})
+    fc.function_call.id = 'another_function_call_id'
+    fc_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=testing_utils.ModelContent([fc]),
+    )
+    fr = Part.from_function_response(
+        name='some_tool', response={'result': 'ok'}
+    )
+    fr.function_response.id = 'test_function_call_id'
+    fr_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=Content(role='user', parts=[fr]),
+    )
+    invocation_context = test_invocation_context([fc_event, fr_event])
+    match = invocation_context._find_matching_function_call(fr_event)
+    assert match is None
+
+  def test_find_matching_function_call_no_call_events(
+      self, test_invocation_context
+  ):
+    """Tests that no matching function call is returned if there are no call events."""
+    fr = Part.from_function_response(
+        name='some_tool', response={'result': 'ok'}
+    )
+    fr.function_response.id = 'test_function_call_id'
+    fr_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=Content(role='user', parts=[fr]),
+    )
+    invocation_context = test_invocation_context([fr_event])
+    match = invocation_context._find_matching_function_call(fr_event)
+    assert match is None
+
+  def test_find_matching_function_call_no_response_in_event(
+      self, test_invocation_context
+  ):
+    """Tests result is None if function_response_event has no function response."""
+    fr_event_no_fr = Event(
+        author='agent',
+        content=Content(role='user', parts=[Part(text='user message')]),
+    )
+    fc = Part.from_function_call(name='some_tool', args={})
+    fc.function_call.id = 'test_function_call_id'
+    fc_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=testing_utils.ModelContent([fc]),
+    )
+    fr = Part.from_function_response(
+        name='some_tool', response={'result': 'ok'}
+    )
+    fr.function_response.id = 'test_function_call_id'
+    fr_event = Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=Content(role='user', parts=[Part(text='user message')]),
+    )
+    invocation_context = test_invocation_context([fc_event, fr_event])
+    match = invocation_context._find_matching_function_call(fr_event_no_fr)
+    assert match is None

tests/unittests/agents/test_resumable_llm_agent.py

@@ -0,0 +1,416 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Union
+
+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.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
+from google.adk.apps.app import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.events.event_actions import EventActions
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.genai.types import Content
+from google.genai.types import Part
+import pytest
+
+from .. import testing_utils
+
+
+def transfer_call_part(agent_name: str) -> Part:
+  return Part.from_function_call(
+      name="transfer_to_agent", args={"agent_name": agent_name}
+  )
+
+
+TRANSFER_RESPONSE_PART = Part.from_function_response(
+    name="transfer_to_agent", response={"result": None}
+)
+
+
+def tool_call_part(tool_name: str) -> Part:
+  part = Part.from_function_call(name=tool_name, args={})
+  part.function_call.id = f"{tool_name}_id"
+  return part
+
+
+def tool_response_part(tool_name: str) -> Part:
+  part = Part.from_function_response(name=tool_name, response={"result": "ok"})
+  part.function_response.id = f"{tool_name}_id"
+  return part
+
+
+def tool_response_part_no_id(tool_name: str) -> Part:
+  part = Part.from_function_response(name=tool_name, response={"result": "ok"})
+  return part
+
+
+END_OF_AGENT = testing_utils.END_OF_AGENT
+
+
+def some_tool():
+  return {"result": "ok"}
+
+
+async def _create_resumable_invocation_context(
+    invocation_id: str, agent: BaseAgent, events: list[Event]
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name="test_app", user_id="test_user"
+  )
+  for event in events:
+    await session_service.append_event(session, event)
+  return InvocationContext(
+      invocation_id=invocation_id,
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=True),
+      run_config=RunConfig(),
+  )
+
+
+async def _resume_and_get_events(
+    agent: BaseAgent, invocation_context: InvocationContext
+) -> list[(str, Union[Part, str])]:
+  events = []
+  async for event in agent.run_async(invocation_context):
+    await invocation_context.session_service.append_event(
+        invocation_context.session, event
+    )
+    events.append(event)
+  return testing_utils.simplify_resumable_app_events(events)
+
+
+class TestResumableLlmAgent:
+  """Test suite for resumable LlmAgent."""
+
+  @pytest.fixture
+  @pytest.mark.asyncio
+  async def resumable_invocation_context(self):
+    """Creates an invocation context for the specified agent."""
+
+    async def factory(agent: BaseAgent, events: list[Event]):
+      return await _create_resumable_invocation_context(
+          invocation_id="test_invocation", agent=agent, events=events
+      )
+
+    return factory
+
+  @pytest.fixture
+  def mock_model(self):
+    """Provides a mock model for the test."""
+
+    def factory(responses: list[Part]):
+      return testing_utils.MockModel.create(responses=responses)
+
+    return factory
+
+  @pytest.mark.asyncio
+  async def test_resume_from_transfer_call(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that the agent resumes from the correct sub-agent after a transfer."""
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=mock_model([
+            "response from sub_agent_1",
+        ]),
+    )
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model(["response from root"]),
+        sub_agents=[sub_agent_1],
+    )
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(
+                parts=[
+                    transfer_call_part("sub_agent_1"),
+                ]
+            ),
+        )
+    ]
+    ctx = await resumable_invocation_context(root_agent, past_events)
+    # Initialize the agent state for the root agent.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("root_agent", TRANSFER_RESPONSE_PART),
+        ("sub_agent_1", "response from sub_agent_1"),
+        ("sub_agent_1", END_OF_AGENT),
+        ("root_agent", END_OF_AGENT),
+    ]
+
+  @pytest.mark.asyncio
+  async def test_resume_from_transfer_response(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that the agent resumes from the correct sub-agent after a transfer."""
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=mock_model([
+            "response from sub_agent_1",
+        ]),
+    )
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model(["response from root"]),
+        sub_agents=[sub_agent_1],
+    )
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(
+                parts=[
+                    TRANSFER_RESPONSE_PART,
+                ]
+            ),
+            actions=EventActions(transfer_to_agent="sub_agent_1"),
+        )
+    ]
+    ctx: InvocationContext = await resumable_invocation_context(
+        root_agent, past_events
+    )
+    # Initialize the agent state for the root agent.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("sub_agent_1", "response from sub_agent_1"),
+        ("sub_agent_1", END_OF_AGENT),
+        ("root_agent", END_OF_AGENT),
+    ]
+
+  @pytest.mark.asyncio
+  async def test_resume_from_model_response(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that no sub-agent is resumed when there has been no transfer."""
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model([
+            "second response from root",
+        ]),
+    )
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[Part(text="initial response from root")]),
+        )
+    ]
+    ctx = await resumable_invocation_context(root_agent, past_events)
+    # Initialize the agent state for the root agent.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("root_agent", "second response from root"),
+        ("root_agent", END_OF_AGENT),
+    ]
+
+  @pytest.mark.asyncio
+  async def test_resume_from_tool_call(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that the agent resumes from a tool call successfully."""
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model(["response after tool call"]),
+        tools=[some_tool],
+    )
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_call_part("some_tool")]),
+        ),
+    ]
+    ctx = await resumable_invocation_context(root_agent, past_events)
+    # Initialize the agent state for the root agent.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("root_agent", tool_response_part_no_id("some_tool")),
+        ("root_agent", "response after tool call"),
+        ("root_agent", END_OF_AGENT),
+    ]
+
+  @pytest.mark.asyncio
+  async def test_resume_after_tool_response(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that the agent does not resume a sub-agent when the user responds to the current agent."""
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model([
+            "response after tool call",
+        ]),
+        tools=[some_tool],
+    )
+
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_call_part("some_tool")]),
+        ),
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_response_part("some_tool")]),
+        ),
+    ]
+    ctx = await resumable_invocation_context(root_agent, past_events)
+    # Initialize the agent state for the root agent.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("root_agent", "response after tool call"),
+        ("root_agent", END_OF_AGENT),
+    ]
+
+  @pytest.mark.asyncio
+  async def test_resume_root_agent_on_user_provided_function_response(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that the agent resumes the correct sub-agent after a user responds to its tool call."""
+
+    def sub_agent_tool():
+      return {"result": "ok"}
+
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=mock_model([
+            "response from sub_agent_1 after tool call",
+        ]),
+        tools=[sub_agent_tool],
+    )
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model(["response from root after tool call"]),
+        sub_agents=[sub_agent_1],
+        tools=[some_tool],
+    )
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            actions=EventActions(transfer_to_agent="sub_agent_1"),
+        ),
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[transfer_call_part("sub_agent_1")]),
+        ),
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[TRANSFER_RESPONSE_PART]),
+            actions=EventActions(transfer_to_agent="sub_agent_1"),
+        ),
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_call_part("some_tool")]),
+        ),
+        Event(
+            author="sub_agent_1",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_call_part("sub_agent_tool")]),
+        ),
+        Event(
+            author="user",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_response_part("some_tool")]),
+        ),
+    ]
+    ctx = await resumable_invocation_context(root_agent, past_events)
+    # Initialize the agent state for the root agent and sub_agent_1.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+    ctx.agent_states[sub_agent_1.name] = BaseAgentState().model_dump(
+        mode="json"
+    )
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("root_agent", "response from root after tool call"),
+        ("root_agent", END_OF_AGENT),
+    ]
+
+  @pytest.mark.asyncio
+  async def test_resume_subagent_on_user_provided_function_response(
+      self, resumable_invocation_context, mock_model
+  ):
+    """Tests that the agent resumes the correct sub-agent after a user responds to its tool call."""
+
+    def sub_agent_tool():
+      return {"result": "ok"}
+
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=mock_model([
+            "response from sub_agent_1 after tool call",
+        ]),
+        tools=[sub_agent_tool],
+    )
+    root_agent = LlmAgent(
+        name="root_agent",
+        model=mock_model(["response from root after tool call"]),
+        sub_agents=[sub_agent_1],
+    )
+    past_events = [
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            actions=EventActions(transfer_to_agent="sub_agent_1"),
+        ),
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[transfer_call_part("sub_agent_1")]),
+        ),
+        Event(
+            author="root_agent",
+            invocation_id="test_invocation",
+            content=Content(parts=[TRANSFER_RESPONSE_PART]),
+            actions=EventActions(transfer_to_agent="sub_agent_1"),
+        ),
+        Event(
+            author="sub_agent_1",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_call_part("sub_agent_tool")]),
+        ),
+        Event(
+            author="user",
+            invocation_id="test_invocation",
+            content=Content(parts=[tool_response_part("sub_agent_tool")]),
+        ),
+    ]
+    ctx = await resumable_invocation_context(root_agent, past_events)
+    # Initialize the agent state for the root agent and sub_agent_1.
+    ctx.agent_states[root_agent.name] = BaseAgentState().model_dump(mode="json")
+    ctx.agent_states[sub_agent_1.name] = BaseAgentState().model_dump(
+        mode="json"
+    )
+
+    assert await _resume_and_get_events(root_agent, ctx) == [
+        ("sub_agent_1", "response from sub_agent_1 after tool call"),
+        ("sub_agent_1", END_OF_AGENT),
+        ("root_agent", END_OF_AGENT),
+    ]

tests/unittests/runners/test_pause_invocation.py

@@ -23,6 +23,7 @@ from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.loop_agent import LoopAgent
 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 App
 from google.adk.apps.app import ResumabilityConfig
 from google.adk.events.event import Event
@@ -71,6 +72,7 @@ class _TestingAgent(BaseAgent):
 _TRANSFER_RESPONSE_PART = Part.from_function_response(
     name="transfer_to_agent", response={"result": None}
 )
+END_OF_AGENT = testing_utils.END_OF_AGENT
 
 
 class BasePauseInvocationTest:
@@ -85,15 +87,15 @@ class BasePauseInvocationTest:
   def app(self, agent: BaseAgent) -> App:
     """Provides an App for the test."""
     return App(
-        name="InMemoryRunner",  # Required for using TestInMemoryRunner.
+        name="test_app",
         root_agent=agent,
         resumability_config=ResumabilityConfig(is_resumable=True),
     )
 
   @pytest.fixture
-  def runner(self, app: App) -> testing_utils.TestInMemoryRunner:
+  def runner(self, app: App) -> testing_utils.InMemoryRunner:
     """Provides an in-memory runner for the agent."""
-    return testing_utils.TestInMemoryRunner(app=app, app_name=None)
+    return testing_utils.InMemoryRunner(app=app)
 
   @staticmethod
   def mock_model(responses: list[Part]) -> testing_utils.MockModel:
@@ -107,10 +109,6 @@ class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
   @pytest.fixture
   def agent(self) -> BaseAgent:
     """Provides a BaseAgent for the test."""
-
-    def test_tool() -> str:
-      return ""
-
     return LlmAgent(
         name="root_agent",
         model=self.mock_model(
@@ -120,14 +118,12 @@ class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_pause_on_long_running_function_call(
+  def test_pause_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a single LlmAgent pauses on long running function call."""
-    assert testing_utils.simplify_events(
+    assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
-        await runner.run_async_with_new_session("test")
-    ) == [
         ("root_agent", Part.from_function_call(name="test_tool", args={})),
     ]
 
@@ -158,34 +154,42 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_pause_first_agent_on_long_running_function_call(
+  def test_pause_first_agent_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
-    """Tests that a single LlmAgent pauses on long running function call."""
+    """Tests that a SequentialAgent pauses on the first sub-agent."""
-    assert testing_utils.simplify_events(
+    assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
-        await runner.run_async_with_new_session("test")
+        (
-    ) == [
+            "root_agent",
+            SequentialAgentState(current_sub_agent="sub_agent_1").model_dump(
+                mode="json"
+            ),
+        ),
         ("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
     ]
 
   @pytest.mark.asyncio
-  async def test_pause_second_agent_on_long_running_function_call(
+  def test_pause_second_agent_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a single LlmAgent pauses on long running function call."""
     # Change the base sequential agent, so that the first agent does not pause.
-    runner.agent.sub_agents[0].tools = [FunctionTool(func=test_tool)]
+    runner.root_agent.sub_agents[0].tools = [FunctionTool(func=test_tool)]
-    runner.agent.sub_agents[0].model = self.mock_model(
+    runner.root_agent.sub_agents[0].model = self.mock_model(
         responses=[
             Part.from_function_call(name="test_tool", args={}),
             Part.from_text(text="model response after tool call"),
         ]
     )
-    assert testing_utils.simplify_events(
+    assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
-        await runner.run_async_with_new_session("test")
+        (
-    ) == [
+            "root_agent",
+            SequentialAgentState(current_sub_agent="sub_agent_1").model_dump(
+                mode="json"
+            ),
+        ),
         ("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
         (
             "sub_agent_1",
@@ -194,6 +198,13 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
             ),
         ),
         ("sub_agent_1", "model response after tool call"),
+        ("sub_agent_1", END_OF_AGENT),
+        (
+            "root_agent",
+            SequentialAgentState(current_sub_agent="sub_agent_2").model_dump(
+                mode="json"
+            ),
+        ),
         ("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
     ]
 
@@ -221,17 +232,19 @@ class TestPauseInvocationWithParallelAgent(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_pause_on_long_running_function_call(
+  def test_pause_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a ParallelAgent pauses on long running function call."""
-    assert testing_utils.simplify_events(
+    simplified_event_parts = testing_utils.simplify_resumable_app_events(
-        await runner.run_async_with_new_session("test")
+        runner.run("test")
-    ) == [
+    )
-        ("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
+    assert (
-        ("sub_agent_2", "Delayed message"),
+        "sub_agent_1",
-    ]
+        Part.from_function_call(name="test_tool", args={}),
+    ) in simplified_event_parts
+    assert ("sub_agent_2", "Delayed message") in simplified_event_parts
 
 
 class TestPauseInvocationWithNestedParallelAgent(BasePauseInvocationTest):
@@ -265,50 +278,49 @@ class TestPauseInvocationWithNestedParallelAgent(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_pause_on_long_running_function_call_in_nested_agent(
+  def test_pause_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a nested ParallelAgent pauses on long running function call."""
-    assert testing_utils.simplify_events(
+    simplified_event_parts = testing_utils.simplify_resumable_app_events(
-        await runner.run_async_with_new_session("test")
+        runner.run("test")
-    ) == [
+    )
-        (
+    assert (
-            "nested_sub_agent_1",
+        "nested_sub_agent_1",
-            Part.from_function_call(name="test_tool", args={}),
+        Part.from_function_call(name="test_tool", args={}),
-        ),
+    ) in simplified_event_parts
-        ("sub_agent_1", "Delayed message"),
+    assert ("sub_agent_1", "Delayed message") in simplified_event_parts
-        ("nested_sub_agent_2", "Delayed message"),
+    assert ("nested_sub_agent_2", "Delayed message") in simplified_event_parts
-    ]
 
   @pytest.mark.asyncio
-  async def test_pause_on_multiple_long_running_function_calls(
+  def test_pause_on_multiple_long_running_function_calls(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a ParallelAgent pauses on long running function calls."""
-    runner.agent.sub_agents[0] = LlmAgent(
+    runner.root_agent.sub_agents[0] = LlmAgent(
         name="sub_agent_1",
         model=self.mock_model(
             responses=[
                 Part.from_function_call(name="test_tool", args={}),
-                Part.from_function_call(name="test_tool", args={}),
             ]
         ),
         tools=[LongRunningFunctionTool(func=test_tool)],
     )
-    simplified_events = testing_utils.simplify_events(
+    simplified_events = testing_utils.simplify_resumable_app_events(
-        await runner.run_async_with_new_session("test")
+        runner.run("test")
     )
-    assert len(simplified_events) == 3
     assert (
         "sub_agent_1",
         Part.from_function_call(name="test_tool", args={}),
     ) in simplified_events
+    assert ("sub_agent_1", END_OF_AGENT) not in simplified_events
     assert (
         "nested_sub_agent_1",
         Part.from_function_call(name="test_tool", args={}),
     ) in simplified_events
+    assert ("nested_sub_agent_1", END_OF_AGENT) not in simplified_events
 
 
 class TestPauseInvocationWithLoopAgent(BasePauseInvocationTest):
@@ -350,15 +362,14 @@ class TestPauseInvocationWithLoopAgent(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_pause_on_long_running_function_call_in_loop(
+  def test_pause_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a LoopAgent pauses on long running function call."""
-    assert testing_utils.simplify_events(
+    assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
-        await runner.run_async_with_new_session("test")
-    ) == [
         ("sub_agent_1", "sub agent 1 response"),
+        ("sub_agent_1", END_OF_AGENT),
         ("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
     ]
 
@@ -400,14 +411,12 @@ class TestPauseInvocationWithLlmAgentTree(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_pause_on_transfer_call_part(
+  def test_pause_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a tree of resumable LlmAgents yields checkpoint events."""
-    assert testing_utils.simplify_events(
+    assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
-        await runner.run_async_with_new_session("test")
-    ) == [
         ("root_agent", _transfer_call_part("sub_llm_agent_1")),
         ("root_agent", _TRANSFER_RESPONSE_PART),
         ("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
@@ -454,14 +463,12 @@ class TestPauseInvocationWithWithTransferLoop(BasePauseInvocationTest):
     )
 
   @pytest.mark.asyncio
-  async def test_agent_tree_yields_checkpoints(
+  def test_pause_on_long_running_function_call(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
   ):
     """Tests that a tree of resumable LlmAgents yields checkpoint events."""
-    assert testing_utils.simplify_events(
+    assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
-        await runner.run_async_with_new_session("test")
-    ) == [
         ("root_agent", _transfer_call_part("sub_llm_agent_1")),
         ("root_agent", _TRANSFER_RESPONSE_PART),
         ("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),

tests/unittests/runners/test_run_tool_confirmation.py

@@ -400,29 +400,31 @@ class TestHITLConfirmationFlowWithResumableApp:
     return LlmAgent(name="root_agent", model=mock_model, tools=tools)
 
   @pytest.fixture
-  def runner(self, agent: LlmAgent) -> testing_utils.TestInMemoryRunner:
+  def runner(self, agent: LlmAgent) -> testing_utils.InMemoryRunner:
     """Provides an in-memory runner for the agent."""
     # Mark the app as resumable. So that the invocation will be paused after the
     # long running tool call.
     app = App(
-        name="InMemoryRunner",  # Required for using TestInMemoryRunner.
+        name="test_app",
         resumability_config=ResumabilityConfig(is_resumable=True),
         root_agent=agent,
     )
-    return testing_utils.TestInMemoryRunner(app=app, app_name=None)
+    return testing_utils.InMemoryRunner(app=app)
 
   @pytest.mark.asyncio
-  async def test_pause_on_request_confirmation(
+  def test_pause_on_request_confirmation(
       self,
-      runner: testing_utils.TestInMemoryRunner,
+      runner: testing_utils.InMemoryRunner,
       agent: LlmAgent,
   ):
     """Tests HITL flow where all tool calls are confirmed."""
-    events = await runner.run_async_with_new_session("test user query")
+    events = runner.run("test user query")
 
     # Verify that the invocation is paused after the long running tool call.
     # So that no intermediate function response and llm response is generated.
-    assert testing_utils.simplify_events(copy.deepcopy(events)) == [
+    assert testing_utils.simplify_resumable_app_events(
+        copy.deepcopy(events)
+    ) == [
         (
             agent.name,
             Part(function_call=FunctionCall(name=agent.tools[0].name, args={})),