chore: import upstream snapshot with attribution

tests/unittests/runners/test_pause_invocation.py

@@ -0,0 +1,538 @@
+# Copyright 2026 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 the resumption flow with different agent structures."""
+
+import asyncio
+from typing import AsyncGenerator
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
+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.apps.app import App
+from google.adk.apps.app import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.tools.exit_loop_tool import exit_loop
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.long_running_tool import LongRunningFunctionTool
+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}
+  )
+
+
+def test_tool():
+  """A test tool; returns None to simulate a pending long-running operation."""
+  return None
+
+
+class _TestingAgent(BaseAgent):
+  """A testing agent that generates an event after a delay."""
+
+  delay: float = 0
+  """The delay before the agent generates an event."""
+
+  def event(self, ctx: InvocationContext):
+    return Event(
+        author=self.name,
+        branch=ctx.branch,
+        invocation_id=ctx.invocation_id,
+        content=testing_utils.ModelContent(
+            parts=[Part.from_text(text="Delayed message")]
+        ),
+    )
+
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    await asyncio.sleep(self.delay)
+    yield self.event(ctx)
+
+
+_TRANSFER_RESPONSE_PART = Part.from_function_response(
+    name="transfer_to_agent", response={"result": None}
+)
+END_OF_AGENT = testing_utils.END_OF_AGENT
+
+
+class BasePauseInvocationTest:
+  """Base class for pausing invocation tests with common fixtures."""
+
+  @pytest.fixture
+  def agent(self) -> BaseAgent:
+    """Provides a BaseAgent for the test."""
+    return BaseAgent(name="test_agent")
+
+  @pytest.fixture
+  def app(self, agent: BaseAgent) -> App:
+    """Provides an App for the test."""
+    return App(
+        name="test_app",
+        root_agent=agent,
+        resumability_config=ResumabilityConfig(is_resumable=True),
+    )
+
+  @pytest.fixture
+  def runner(self, app: App) -> testing_utils.InMemoryRunner:
+    """Provides an in-memory runner for the agent."""
+    return testing_utils.InMemoryRunner(app=app)
+
+  @staticmethod
+  def mock_model(responses: list[Part]) -> testing_utils.MockModel:
+    """Provides a mock model with predefined responses."""
+    return testing_utils.MockModel.create(responses=responses)
+
+
+class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
+  """Tests the resumption flow with a single LlmAgent."""
+
+  @pytest.fixture
+  def agent(self) -> BaseAgent:
+    """Provides a BaseAgent for the test."""
+    return LlmAgent(
+        name="root_agent",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+
+  @pytest.mark.asyncio
+  def test_pause_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a single LlmAgent pauses on long running function call."""
+    actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
+    behavioral = [e for e in actual if not isinstance(e[1], dict)]
+    assert behavioral == [
+        # execute_tools yields the interrupt event with long_running_tool_ids.
+        ("root_agent", Part.from_function_call(name="test_tool", args={})),
+    ]
+
+
+class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
+  """Tests pausing invocation with a SequentialAgent."""
+
+  @pytest.fixture
+  def agent(self) -> BaseAgent:
+    """Provides a BaseAgent for the test."""
+    sub_agent1 = LlmAgent(
+        name="sub_agent_1",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    sub_agent2 = LlmAgent(
+        name="sub_agent_2",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    return SequentialAgent(
+        name="root_agent",
+        sub_agents=[sub_agent1, sub_agent2],
+    )
+
+  @pytest.mark.asyncio
+  def test_pause_first_agent_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a SequentialAgent pauses on the first sub-agent."""
+    actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
+    behavioral = [e for e in actual if not isinstance(e[1], dict)]
+    assert behavioral == [
+        ("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
+    ]
+
+  @pytest.mark.xfail(
+      reason=(
+          "Tests implementation details that are different in V2 and will be"
+          " deprecated."
+      )
+  )
+  @pytest.mark.asyncio
+  def test_pause_second_agent_on_long_running_function_call(
+      self,
+  ):
+    """Tests that a single LlmAgent pauses on long running function call."""
+    # Construct sub_agent_1 with regular FunctionTool (not long-running).
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=self.mock_model(
+            responses=[
+                Part.from_function_call(name="test_tool", args={}),
+                Part.from_text(text="model response after tool call"),
+            ]
+        ),
+        tools=[FunctionTool(func=test_tool)],
+    )
+    sub_agent_2 = LlmAgent(
+        name="sub_agent_2",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    agent = SequentialAgent(
+        name="root_agent",
+        sub_agents=[sub_agent_1, sub_agent_2],
+    )
+    app = App(
+        name="test_app",
+        root_agent=agent,
+        resumability_config=ResumabilityConfig(is_resumable=True),
+    )
+    runner = testing_utils.InMemoryRunner(app=app)
+    actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
+    behavioral = [e for e in actual if not isinstance(e[1], dict)]
+    assert behavioral == [
+        ("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
+        (
+            "sub_agent_1",
+            Part.from_function_response(
+                name="test_tool", response={"result": None}
+            ),
+        ),
+        ("sub_agent_1", "model response after tool call"),
+        # Wrapper emits output before END_OF_AGENT.
+        ("root_agent", "model response after tool call"),
+        ("sub_agent_1", END_OF_AGENT),
+        ("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
+    ]
+
+
+class TestPauseInvocationWithParallelAgent(BasePauseInvocationTest):
+  """Tests pausing invocation with a ParallelAgent."""
+
+  @pytest.fixture
+  def agent(self) -> BaseAgent:
+    """Provides a BaseAgent for the test."""
+    sub_agent1 = LlmAgent(
+        name="sub_agent_1",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    sub_agent2 = _TestingAgent(
+        name="sub_agent_2",
+        delay=0.5,
+    )
+    return ParallelAgent(
+        name="root_agent",
+        sub_agents=[sub_agent1, sub_agent2],
+    )
+
+  @pytest.mark.asyncio
+  def test_pause_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a ParallelAgent pauses on long running function call."""
+    simplified_event_parts = testing_utils.simplify_resumable_app_events(
+        runner.run("test")
+    )
+    assert (
+        "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):
+  """Tests pausing invocation with a nested ParallelAgent."""
+
+  @pytest.fixture
+  def agent(self) -> BaseAgent:
+    """Provides a BaseAgent for the test."""
+    nested_sub_agent_1 = LlmAgent(
+        name="nested_sub_agent_1",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    nested_sub_agent_2 = _TestingAgent(
+        name="nested_sub_agent_2",
+        delay=0.5,
+    )
+    nested_parallel_agent = ParallelAgent(
+        name="nested_parallel_agent",
+        sub_agents=[nested_sub_agent_1, nested_sub_agent_2],
+    )
+    sub_agent_1 = _TestingAgent(
+        name="sub_agent_1",
+        delay=0.5,
+    )
+    return ParallelAgent(
+        name="root_agent",
+        sub_agents=[sub_agent_1, nested_parallel_agent],
+    )
+
+  @pytest.mark.asyncio
+  def test_pause_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a nested ParallelAgent pauses on long running function call."""
+    simplified_event_parts = testing_utils.simplify_resumable_app_events(
+        runner.run("test")
+    )
+    assert (
+        "nested_sub_agent_1",
+        Part.from_function_call(name="test_tool", args={}),
+    ) in simplified_event_parts
+    assert ("sub_agent_1", "Delayed message") in simplified_event_parts
+    assert ("nested_sub_agent_2", "Delayed message") in simplified_event_parts
+
+  @pytest.mark.asyncio
+  def test_pause_on_multiple_long_running_function_calls(
+      self,
+  ):
+    """Tests that a ParallelAgent pauses on long running function calls."""
+    nested_sub_agent_1 = LlmAgent(
+        name="nested_sub_agent_1",
+        model=self.mock_model(
+            responses=[Part.from_function_call(name="test_tool", args={})]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    nested_sub_agent_2 = _TestingAgent(
+        name="nested_sub_agent_2",
+        delay=0.5,
+    )
+    nested_parallel_agent = ParallelAgent(
+        name="nested_parallel_agent",
+        sub_agents=[nested_sub_agent_1, nested_sub_agent_2],
+    )
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=self.mock_model(
+            responses=[
+                Part.from_function_call(name="test_tool", args={}),
+            ]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    agent = ParallelAgent(
+        name="root_agent",
+        sub_agents=[sub_agent_1, nested_parallel_agent],
+    )
+    app = App(
+        name="test_app",
+        root_agent=agent,
+        resumability_config=ResumabilityConfig(is_resumable=True),
+    )
+    runner = testing_utils.InMemoryRunner(app=app)
+    simplified_events = testing_utils.simplify_resumable_app_events(
+        runner.run("test")
+    )
+    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):
+  """Tests pausing invocation with a LoopAgent."""
+
+  @pytest.fixture
+  def agent(self) -> BaseAgent:
+    """Provides a BaseAgent for the test."""
+    sub_agent_1 = LlmAgent(
+        name="sub_agent_1",
+        model=self.mock_model(
+            responses=[
+                Part.from_text(text="sub agent 1 response"),
+            ]
+        ),
+    )
+    sub_agent_2 = LlmAgent(
+        name="sub_agent_2",
+        model=self.mock_model(
+            responses=[
+                Part.from_function_call(name="test_tool", args={}),
+            ]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    sub_agent_3 = LlmAgent(
+        name="sub_agent_3",
+        model=self.mock_model(
+            responses=[
+                Part.from_function_call(name="exit_loop", args={}),
+            ]
+        ),
+        tools=[exit_loop],
+    )
+    return LoopAgent(
+        name="root_agent",
+        sub_agents=[sub_agent_1, sub_agent_2, sub_agent_3],
+        max_iterations=2,
+    )
+
+  @pytest.mark.xfail(
+      reason=(
+          "Tests implementation details that are different in V2 and will be"
+          " deprecated."
+      )
+  )
+  @pytest.mark.asyncio
+  def test_pause_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a LoopAgent pauses on long running function call."""
+    actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
+    behavioral = [e for e in actual if not isinstance(e[1], dict)]
+    assert behavioral == [
+        ("sub_agent_1", "sub agent 1 response"),
+        # Wrapper emits output before END_OF_AGENT.
+        ("root_agent", "sub agent 1 response"),
+        ("sub_agent_1", END_OF_AGENT),
+        ("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
+    ]
+
+
+class TestPauseInvocationWithLlmAgentTree(BasePauseInvocationTest):
+  """Tests the pausing invocation with a tree of LlmAgents."""
+
+  @pytest.fixture
+  def agent(self) -> LlmAgent:
+    """Provides an LlmAgent with sub-agents for the test."""
+    sub_llm_agent_1 = LlmAgent(
+        name="sub_llm_agent_1",
+        model=self.mock_model(
+            responses=[
+                _transfer_call_part("sub_llm_agent_2"),
+                "llm response not used",
+            ]
+        ),
+    )
+    sub_llm_agent_2 = LlmAgent(
+        name="sub_llm_agent_2",
+        model=self.mock_model(
+            responses=[
+                Part.from_function_call(name="test_tool", args={}),
+                "llm response not used",
+            ]
+        ),
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+    return LlmAgent(
+        name="root_agent",
+        model=self.mock_model(
+            responses=[
+                _transfer_call_part("sub_llm_agent_1"),
+                "llm response not used",
+            ]
+        ),
+        sub_agents=[sub_llm_agent_1, sub_llm_agent_2],
+    )
+
+  @pytest.mark.asyncio
+  def test_pause_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a tree of resumable LlmAgents yields checkpoint events."""
+    actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
+    behavioral = [e for e in actual if not isinstance(e[1], dict)]
+    assert behavioral == [
+        ("root_agent", _transfer_call_part("sub_llm_agent_1")),
+        ("root_agent", _TRANSFER_RESPONSE_PART),
+        ("root_agent", END_OF_AGENT),
+        ("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
+        ("sub_llm_agent_1", _TRANSFER_RESPONSE_PART),
+        ("sub_llm_agent_1", END_OF_AGENT),
+        ("sub_llm_agent_2", Part.from_function_call(name="test_tool", args={})),
+    ]
+
+
+class TestPauseInvocationWithWithTransferLoop(BasePauseInvocationTest):
+  """Tests pausing the invocation when the agent transfer forms a loop."""
+
+  @pytest.fixture
+  def agent(self) -> LlmAgent:
+    """Provides an LlmAgent with sub-agents for the test."""
+    sub_llm_agent_1 = LlmAgent(
+        name="sub_llm_agent_1",
+        model=self.mock_model(
+            responses=[
+                _transfer_call_part("sub_llm_agent_2"),
+                "llm response not used",
+            ]
+        ),
+    )
+    sub_llm_agent_2 = LlmAgent(
+        name="sub_llm_agent_2",
+        model=self.mock_model(
+            responses=[
+                _transfer_call_part("root_agent"),
+                "llm response not used",
+            ]
+        ),
+    )
+    return LlmAgent(
+        name="root_agent",
+        model=self.mock_model(
+            responses=[
+                _transfer_call_part("sub_llm_agent_1"),
+                Part.from_function_call(name="test_tool", args={}),
+                "llm response not used",
+            ]
+        ),
+        sub_agents=[sub_llm_agent_1, sub_llm_agent_2],
+        tools=[LongRunningFunctionTool(func=test_tool)],
+    )
+
+  @pytest.mark.asyncio
+  def test_pause_on_long_running_function_call(
+      self,
+      runner: testing_utils.InMemoryRunner,
+  ):
+    """Tests that a tree of resumable LlmAgents yields checkpoint events."""
+    actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
+    behavioral = [e for e in actual if not isinstance(e[1], dict)]
+    assert behavioral == [
+        ("root_agent", _transfer_call_part("sub_llm_agent_1")),
+        ("root_agent", _TRANSFER_RESPONSE_PART),
+        ("root_agent", END_OF_AGENT),
+        ("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
+        ("sub_llm_agent_1", _TRANSFER_RESPONSE_PART),
+        ("sub_llm_agent_1", END_OF_AGENT),
+        ("sub_llm_agent_2", _transfer_call_part("root_agent")),
+        ("sub_llm_agent_2", _TRANSFER_RESPONSE_PART),
+        ("sub_llm_agent_2", END_OF_AGENT),
+        ("root_agent", Part.from_function_call(name="test_tool", args={})),
+    ]