feat: Support resuming from a paused invocation starting from a sub-agent

PiperOrigin-RevId: 817766247

Author: XinranTang

src/google/adk/runners.py

@@ -381,12 +381,11 @@ async def _run_with_trace(
               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]
+          if invocation_context.end_of_agents.get(
+              invocation_context.agent.name
+          ):
+            # Directly return if the current agent in invocation context is
+            # already final.
             return
         else:
           invocation_context = await self._setup_context_for_new_invocation(
@@ -869,6 +868,13 @@ async def _setup_context_for_resumed_invocation(
       )
     # Step 4: Populate agent states for the current invocation.
     invocation_context.populate_invocation_agent_states()
+    # Step 5: Set agent to run for the invocation.
+    #
+    # If the root agent is not found in end_of_agents, it means the invocation
+    # started from a sub-agent and paused on a sub-agent.
+    # We should find the appropriate agent to run to continue the invocation.
+    if self.agent.name not in invocation_context.end_of_agents:
+      invocation_context.agent = self._find_agent_to_run(session, self.agent)
     return invocation_context
 
   def _find_user_message_for_invocation(

tests/unittests/runners/test_resume_invocation.py

@@ -0,0 +1,146 @@
+# 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.
+"""Tests for edge cases of resuming invocations."""
+
+import copy
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.apps.app import App
+from google.adk.apps.app import ResumabilityConfig
+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}
+)
+
+
+@pytest.mark.asyncio
+async def test_resume_invocation_from_sub_agent():
+  """A test case for an edge case, where an invocation-to-resume starts from a sub-agent.
+
+  For example:
+    invocation1: root_agent -> sub_agent
+    invocation2: sub_agent [paused][resume]
+  """
+  # Step 1: Setup
+  # root_agent -> sub_agent
+  sub_agent = LlmAgent(
+      name="sub_agent",
+      model=testing_utils.MockModel.create(
+          responses=[
+              "first response from sub_agent",
+              "second response from sub_agent",
+              "third response from sub_agent",
+          ]
+      ),
+  )
+  root_agent = LlmAgent(
+      name="root_agent",
+      model=testing_utils.MockModel.create(
+          responses=[transfer_call_part(sub_agent.name)]
+      ),
+      sub_agents=[sub_agent],
+  )
+  runner = testing_utils.InMemoryRunner(
+      app=App(
+          name="test_app",
+          root_agent=root_agent,
+          resumability_config=ResumabilityConfig(is_resumable=True),
+      )
+  )
+
+  # Step 2: Run the first invocation
+  # Expect the invocation to start from root_agent and transferred to sub_agent.
+  invocation_1_events = runner.run("test user query")
+  assert testing_utils.simplify_resumable_app_events(
+      copy.deepcopy(invocation_1_events)
+  ) == [
+      (
+          root_agent.name,
+          transfer_call_part(sub_agent.name),
+      ),
+      (
+          root_agent.name,
+          TRANSFER_RESPONSE_PART,
+      ),
+      (
+          sub_agent.name,
+          "first response from sub_agent",
+      ),
+      (
+          sub_agent.name,
+          testing_utils.END_OF_AGENT,
+      ),
+      (
+          root_agent.name,
+          testing_utils.END_OF_AGENT,
+      ),
+  ]
+
+  # Step 3: Run the second invocation
+  # Expect the invocation to directly start from sub_agent.
+  invocation_2_events = runner.run(
+      "test user query 2",
+  )
+  assert testing_utils.simplify_resumable_app_events(
+      copy.deepcopy(invocation_2_events)
+  ) == [
+      (
+          sub_agent.name,
+          "second response from sub_agent",
+      ),
+      (sub_agent.name, testing_utils.END_OF_AGENT),
+  ]
+  # Asserts the invocation will be a no-op if the current agent in context is
+  # already final.
+  assert not await runner.run_async(
+      invocation_id=invocation_2_events[0].invocation_id
+  )
+
+  # Step 4: Copy all session.events[:-1] to a new session
+  # This is to simulate the case where we pause on the second invocation.
+  session_id = runner.session_id
+  session = await runner.runner.session_service.get_session(
+      app_name="test_app", user_id="test_user", session_id=session_id
+  )
+  new_session = await runner.runner.session_service.create_session(
+      app_name=session.app_name, user_id=session.user_id
+  )
+  for event in session.events[:-1]:
+    await runner.runner.session_service.append_event(new_session, event)
+  runner.session_id = new_session.id
+
+  # Step 5: Resume the second invocation
+  resumed_invocation_2_events = await runner.run_async(
+      invocation_id=invocation_2_events[0].invocation_id
+  )
+  assert testing_utils.simplify_resumable_app_events(
+      copy.deepcopy(resumed_invocation_2_events)
+  ) == [
+      (
+          sub_agent.name,
+          "third response from sub_agent",
+      ),
+      (sub_agent.name, testing_utils.END_OF_AGENT),
+  ]

tests/unittests/testing_utils.py

@@ -274,14 +274,16 @@ def run(self, new_message: types.ContentUnion) -> list[Event]:
     )
 
   async def run_async(
-      self, new_message: types.ContentUnion, invocation_id: Optional[str] = None
+      self,
+      new_message: Optional[types.ContentUnion] = None,
+      invocation_id: Optional[str] = None,
   ) -> list[Event]:
     events = []
     async for event in self.runner.run_async(
         user_id=self.session.user_id,
         session_id=self.session.id,
         invocation_id=invocation_id,
-        new_message=get_user_content(new_message),
+        new_message=get_user_content(new_message) if new_message else None,
     ):
       events.append(event)
     return events