chore: import upstream snapshot with attribution

tests/unittests/agents/test_resumable_llm_agent.py

@@ -0,0 +1,264 @@
+# 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 resumable LlmAgent scenarios.
+
+Verifies that the Mesh-based LlmAgent correctly resumes from various
+states: after transfers, tool calls, tool responses, and with
+sub-agent tool calls.
+"""
+
+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}
+)
+
+END_OF_AGENT = testing_utils.END_OF_AGENT
+
+
+def some_tool():
+  return {"result": "ok"}
+
+
+def _behavioral_events(events):
+  """Extract behavioral events (non-state) from resumable app events."""
+  return [
+      e
+      for e in testing_utils.simplify_resumable_app_events(
+          copy.deepcopy(events)
+      )
+      if not isinstance(e[1], dict)
+  ]
+
+
+@pytest.mark.asyncio
+async def test_resume_from_transfer():
+  """Tests that the agent resumes from the correct sub-agent after a transfer.
+
+  invocation1: root_agent transfers to sub_agent_1
+  invocation2: sub_agent_1 responds (resuming from the transfer)
+  """
+  sub_agent_1 = LlmAgent(
+      name="sub_agent_1",
+      model=testing_utils.MockModel.create(
+          responses=[
+              "response from sub_agent_1",
+              "second response from sub_agent_1",
+          ]
+      ),
+  )
+  root_agent = LlmAgent(
+      name="root_agent",
+      model=testing_utils.MockModel.create(
+          responses=[transfer_call_part("sub_agent_1")]
+      ),
+      sub_agents=[sub_agent_1],
+  )
+  runner = testing_utils.InMemoryRunner(
+      app=App(
+          name="test_app",
+          root_agent=root_agent,
+          resumability_config=ResumabilityConfig(is_resumable=True),
+      )
+  )
+
+  # Invocation 1: root transfers to sub_agent_1.
+  inv1_events = await runner.run_async("test query")
+  inv1_behavioral = _behavioral_events(inv1_events)
+  assert inv1_behavioral == [
+      ("root_agent", transfer_call_part("sub_agent_1")),
+      ("root_agent", TRANSFER_RESPONSE_PART),
+      ("root_agent", END_OF_AGENT),
+      ("sub_agent_1", "response from sub_agent_1"),
+      ("sub_agent_1", END_OF_AGENT),
+  ]
+
+  # Invocation 2: sub_agent_1 is now active, should respond directly.
+  inv2_events = await runner.run_async("follow up query")
+  inv2_behavioral = _behavioral_events(inv2_events)
+  assert inv2_behavioral == [
+      ("sub_agent_1", "second response from sub_agent_1"),
+      ("sub_agent_1", END_OF_AGENT),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_resume_from_model_response():
+  """Tests that the root agent resumes when there has been no transfer."""
+  root_agent = LlmAgent(
+      name="root_agent",
+      model=testing_utils.MockModel.create(
+          responses=[
+              "first response from root",
+              "second response from root",
+          ]
+      ),
+  )
+  runner = testing_utils.InMemoryRunner(
+      app=App(
+          name="test_app",
+          root_agent=root_agent,
+          resumability_config=ResumabilityConfig(is_resumable=True),
+      )
+  )
+
+  # Invocation 1: root responds normally.
+  inv1_events = await runner.run_async("test query")
+  inv1_behavioral = _behavioral_events(inv1_events)
+  assert inv1_behavioral == [
+      ("root_agent", "first response from root"),
+      ("root_agent", END_OF_AGENT),
+  ]
+
+  # Invocation 2: root should respond again (no transfer happened).
+  inv2_events = await runner.run_async("follow up")
+  inv2_behavioral = _behavioral_events(inv2_events)
+  assert inv2_behavioral == [
+      ("root_agent", "second response from root"),
+      ("root_agent", END_OF_AGENT),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_resume_from_tool_call():
+  """Tests that the agent resumes from a tool call.
+
+  invocation1: root_agent calls some_tool, gets response, then responds
+  invocation2: root_agent responds again (tool was non-long-running)
+  """
+  root_agent = LlmAgent(
+      name="root_agent",
+      model=testing_utils.MockModel.create(
+          responses=[
+              Part.from_function_call(name="some_tool", args={}),
+              "response after tool call",
+              "second response",
+          ]
+      ),
+      tools=[some_tool],
+  )
+  runner = testing_utils.InMemoryRunner(
+      app=App(
+          name="test_app",
+          root_agent=root_agent,
+          resumability_config=ResumabilityConfig(is_resumable=True),
+      )
+  )
+
+  # Invocation 1: root calls tool, gets response, responds.
+  inv1_events = await runner.run_async("test query")
+  inv1_behavioral = _behavioral_events(inv1_events)
+  assert inv1_behavioral == [
+      ("root_agent", Part.from_function_call(name="some_tool", args={})),
+      (
+          "root_agent",
+          Part.from_function_response(
+              name="some_tool", response={"result": "ok"}
+          ),
+      ),
+      ("root_agent", "response after tool call"),
+      ("root_agent", END_OF_AGENT),
+  ]
+
+  # Invocation 2: root resumes normally.
+  inv2_events = await runner.run_async("follow up")
+  inv2_behavioral = _behavioral_events(inv2_events)
+  assert inv2_behavioral == [
+      ("root_agent", "second response"),
+      ("root_agent", END_OF_AGENT),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_resume_subagent_after_transfer_and_tool_call():
+  """Tests resuming a sub-agent that called a tool after being transferred to.
+
+  invocation1: root_agent transfers to sub_agent_1, sub_agent_1 calls tool
+               and responds
+  invocation2: sub_agent_1 is still active, responds directly
+  """
+
+  def sub_agent_tool():
+    return {"result": "ok"}
+
+  sub_agent_1 = LlmAgent(
+      name="sub_agent_1",
+      model=testing_utils.MockModel.create(
+          responses=[
+              Part.from_function_call(name="sub_agent_tool", args={}),
+              "response from sub_agent_1 after tool",
+              "second response from sub_agent_1",
+          ]
+      ),
+      tools=[sub_agent_tool],
+  )
+  root_agent = LlmAgent(
+      name="root_agent",
+      model=testing_utils.MockModel.create(
+          responses=[transfer_call_part("sub_agent_1")]
+      ),
+      sub_agents=[sub_agent_1],
+  )
+  runner = testing_utils.InMemoryRunner(
+      app=App(
+          name="test_app",
+          root_agent=root_agent,
+          resumability_config=ResumabilityConfig(is_resumable=True),
+      )
+  )
+
+  # Invocation 1: root transfers, sub_agent calls tool and responds.
+  inv1_events = await runner.run_async("test query")
+  inv1_behavioral = _behavioral_events(inv1_events)
+  assert inv1_behavioral == [
+      ("root_agent", transfer_call_part("sub_agent_1")),
+      ("root_agent", TRANSFER_RESPONSE_PART),
+      ("root_agent", END_OF_AGENT),
+      (
+          "sub_agent_1",
+          Part.from_function_call(name="sub_agent_tool", args={}),
+      ),
+      (
+          "sub_agent_1",
+          Part.from_function_response(
+              name="sub_agent_tool", response={"result": "ok"}
+          ),
+      ),
+      ("sub_agent_1", "response from sub_agent_1 after tool"),
+      ("sub_agent_1", END_OF_AGENT),
+  ]
+
+  # Invocation 2: sub_agent_1 is still active, responds directly.
+  inv2_events = await runner.run_async("follow up")
+  inv2_behavioral = _behavioral_events(inv2_events)
+  assert inv2_behavioral == [
+      ("sub_agent_1", "second response from sub_agent_1"),
+      ("sub_agent_1", END_OF_AGENT),
+  ]