chore: import upstream snapshot with attribution

tests/unittests/agents/__init__.py

@@ -0,0 +1,13 @@
+# 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/unittests/agents/llm/__init__.py

@@ -0,0 +1,13 @@
+# 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/unittests/agents/llm/event_utils.py

@@ -0,0 +1,81 @@
+# 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.
+
+"""Shared test helpers for extracting data from Event lists."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from google.adk.events.event import Event
+
+
+async def collect_events(node, ctx, node_input=None) -> list[Event]:
+  """Collect all events yielded by node.run()."""
+  events = []
+  async for event in node.run(ctx=ctx, node_input=node_input):
+    events.append(event)
+  return events
+
+
+def output_events(events: list[Event]) -> list[Event]:
+  """Filter to events that carry output."""
+  return [e for e in events if e.output is not None]
+
+
+def content_events(events: list[Event]) -> list[Event]:
+  """Filter to events that have content."""
+  return [e for e in events if e.content and e.content.parts]
+
+
+def text_parts(events: list[Event]) -> list[str]:
+  """Extract text strings from content events."""
+  return [
+      p.text.strip()
+      for e in content_events(events)
+      for p in e.content.parts
+      if p.text
+  ]
+
+
+def function_call_names(events: list[Event]) -> list[str]:
+  """Extract function call names from content events."""
+  return [
+      p.function_call.name
+      for e in content_events(events)
+      for p in e.content.parts
+      if p.function_call
+  ]
+
+
+def function_response_dicts(events: list[Event]) -> list[dict[str, Any]]:
+  """Extract function response dicts from events."""
+  return [
+      dict(p.function_response.response or {})
+      for e in content_events(events)
+      for p in e.content.parts
+      if p.function_response
+  ]
+
+
+def function_responses_by_name(
+    events: list[Event],
+) -> dict[str, dict[str, Any]]:
+  """Extract {tool_name: response_dict} from function response events."""
+  return {
+      p.function_response.name: dict(p.function_response.response or {})
+      for e in content_events(events)
+      for p in e.content.parts
+      if p.function_response
+  }

tests/unittests/agents/llm/task/__init__.py

@@ -0,0 +1,13 @@
+# 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/unittests/agents/llm/task/test_finish_task_tool.py

@@ -0,0 +1,363 @@
+# 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 finish_task_tool module."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest import mock
+
+from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME
+from google.adk.agents.llm.task._finish_task_tool import FinishTaskTool
+from pydantic import BaseModel
+import pytest
+
+
+class SampleOutputSchema(BaseModel):
+  """Sample Pydantic model for testing output schema."""
+
+  result: str
+  count: int
+
+
+class NestedOutputSchema(BaseModel):
+  """Nested Pydantic model for testing complex schemas."""
+
+  name: str
+  details: dict
+
+
+def _make_task_agent(
+    name='test_agent',
+    mode='task',
+    output_schema=None,
+):
+  """Create a mock task agent for FinishTaskTool construction."""
+  agent = mock.MagicMock()
+  agent.name = name
+  agent.mode = mode
+  agent.output_schema = output_schema
+  return agent
+
+
+class TestFinishTaskTool:
+  """Tests for the FinishTaskTool class."""
+
+  def test_init_without_output_schema(self):
+    """Test FinishTaskTool initialization without output schema."""
+    agent = _make_task_agent()
+    tool = FinishTaskTool(task_agent=agent)
+    assert tool.name == FINISH_TASK_TOOL_NAME
+    assert 'Signal that this agent has completed' in tool.description
+    assert 'output data' not in tool.description
+
+  def test_init_with_output_schema(self):
+    """Test FinishTaskTool initialization with output schema."""
+    agent = _make_task_agent(output_schema=SampleOutputSchema)
+    tool = FinishTaskTool(task_agent=agent)
+    assert tool.name == FINISH_TASK_TOOL_NAME
+    assert tool.output_schema is SampleOutputSchema
+    assert 'Signal that this agent has completed' in tool.description
+    assert 'output data' in tool.description
+
+  def test_get_declaration_without_output_schema(self):
+    """Test function declaration generation without output schema."""
+    agent = _make_task_agent()
+    tool = FinishTaskTool(task_agent=agent)
+    declaration = tool._get_declaration()
+
+    assert declaration is not None
+    assert declaration.name == FINISH_TASK_TOOL_NAME
+    assert declaration.parameters_json_schema is not None
+    schema = declaration.parameters_json_schema
+    assert 'result' in schema.get('properties', {})
+
+  def test_get_declaration_with_output_schema(self):
+    """Test function declaration generation with output schema."""
+    agent = _make_task_agent(output_schema=SampleOutputSchema)
+    tool = FinishTaskTool(task_agent=agent)
+    declaration = tool._get_declaration()
+
+    assert declaration is not None
+    assert declaration.name == FINISH_TASK_TOOL_NAME
+    assert declaration.parameters_json_schema is not None
+    schema = declaration.parameters_json_schema
+    assert 'result' in schema.get('properties', {})
+    assert 'count' in schema.get('properties', {})
+
+  @pytest.mark.asyncio
+  async def test_run_async_returns_confirmation(self):
+    """Test that run_async returns a confirmation message."""
+    agent = _make_task_agent()
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+
+    result = await tool.run_async(
+        args={'result': 'done'}, tool_context=mock_tool_context
+    )
+
+    assert result == 'Task completed.'
+
+  @pytest.mark.asyncio
+  async def test_run_async_with_args(self):
+    """Test that run_async validates args and returns confirmation."""
+    agent = _make_task_agent(output_schema=SampleOutputSchema)
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+
+    result = await tool.run_async(
+        args={'result': 'success', 'count': 42},
+        tool_context=mock_tool_context,
+    )
+
+    assert result == 'Task completed.'
+
+
+class TestBuildInstruction:
+  """Tests for the _build_instruction method."""
+
+  def test_instruction_content(self):
+    """Test instruction generation contains expected content."""
+    agent = _make_task_agent()
+    tool = FinishTaskTool(task_agent=agent)
+    instruction = tool._build_instruction()
+
+    assert 'finish_task' in instruction
+    assert 'Do NOT call `finish_task` prematurely' in instruction
+    assert 'call `finish_task` by itself with' in instruction
+
+
+class TestProcessLlmRequest:
+  """Tests for the process_llm_request method."""
+
+  @pytest.mark.asyncio
+  async def test_process_llm_request_adds_tool_and_instruction(self):
+    """Test that process_llm_request adds tool and instruction."""
+    agent = _make_task_agent()
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+    mock_tool_context._invocation_context.branch = None
+    mock_tool_context.session.events = []
+    mock_llm_request = mock.MagicMock()
+    mock_llm_request.append_tools = mock.MagicMock()
+    mock_llm_request.append_instructions = mock.MagicMock()
+
+    await tool.process_llm_request(
+        tool_context=mock_tool_context, llm_request=mock_llm_request
+    )
+
+    # Should add tool via parent's process_llm_request
+    mock_llm_request.append_tools.assert_called_once_with([tool])
+    # Should append instruction (at least the base instruction)
+    mock_llm_request.append_instructions.assert_called()
+    instruction_arg = mock_llm_request.append_instructions.call_args_list[0][0][
+        0
+    ]
+    assert len(instruction_arg) == 1
+    assert 'finish_task' in instruction_arg[0]
+
+
+class TestFinishTaskToolName:
+  """Tests for the FINISH_TASK_TOOL_NAME constant."""
+
+  def test_constant_value(self):
+    """Test that the constant has the expected value."""
+    assert FINISH_TASK_TOOL_NAME == 'finish_task'
+
+
+class TestFinishTaskToolValidation:
+  """Tests for the FinishTaskTool argument validation."""
+
+  @pytest.mark.asyncio
+  async def test_run_async_validation_error_missing_required_field(self):
+    """Test that validation error is returned when required fields are missing."""
+    agent = _make_task_agent(output_schema=SampleOutputSchema)
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+
+    # Missing 'count' field which is required
+    result = await tool.run_async(
+        args={'result': 'success'},
+        tool_context=mock_tool_context,
+    )
+
+    assert isinstance(result, dict)
+    assert 'error' in result
+    assert 'finish_task' in result['error']
+    assert 'validation errors' in result['error']
+    assert 'count' in result['error']
+
+  @pytest.mark.asyncio
+  async def test_run_async_validation_error_wrong_type(self):
+    """Test that validation error is returned when types are wrong."""
+    agent = _make_task_agent(output_schema=SampleOutputSchema)
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+
+    # 'count' should be int, not string
+    result = await tool.run_async(
+        args={'result': 'success', 'count': 'not_an_int'},
+        tool_context=mock_tool_context,
+    )
+
+    assert isinstance(result, dict)
+    assert 'error' in result
+    assert 'validation errors' in result['error']
+
+  @pytest.mark.asyncio
+  async def test_run_async_validation_passes_with_valid_args(self):
+    """Test that validation passes with valid args."""
+    agent = _make_task_agent(output_schema=SampleOutputSchema)
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+
+    result = await tool.run_async(
+        args={'result': 'success', 'count': 42},
+        tool_context=mock_tool_context,
+    )
+
+    assert result == 'Task completed.'
+
+
+class TestFinishTaskToolAllSchemaTypes:
+  """Tests for FinishTaskTool across all supported SchemaType variants.
+
+  Object schemas (BaseModel, dict) use parameters directly.
+  Non-object schemas (str, int, bool, float, list) are wrapped in a
+  JSON object with a 'result' key so they can be used as function
+  parameters.
+  """
+
+  @pytest.mark.parametrize(
+      'output_schema, expected_wrapper_key',
+      [
+          (SampleOutputSchema, None),
+          (dict[str, Any], None),
+          (str, 'result'),
+          (int, 'result'),
+          (bool, 'result'),
+          (float, 'result'),
+          (list[str], 'result'),
+          (list[int], 'result'),
+          (list[SampleOutputSchema], 'result'),
+      ],
+      ids=[
+          'BaseModel',
+          'dict',
+          'str',
+          'int',
+          'bool',
+          'float',
+          'list_str',
+          'list_int',
+          'list_BaseModel',
+      ],
+  )
+  def test_wrapper_key(self, output_schema, expected_wrapper_key):
+    """Verify wrapper key is set correctly for each schema type."""
+    agent = _make_task_agent(output_schema=output_schema)
+    tool = FinishTaskTool(task_agent=agent)
+    assert tool._wrapper_key == expected_wrapper_key
+
+  @pytest.mark.parametrize(
+      'output_schema',
+      [str, int, bool, float, list[str], list[int], list[SampleOutputSchema]],
+      ids=[
+          'str',
+          'int',
+          'bool',
+          'float',
+          'list_str',
+          'list_int',
+          'list_BaseModel',
+      ],
+  )
+  def test_get_declaration_wrapped_schema(self, output_schema):
+    """Non-object schemas produce a declaration with 'result' property."""
+    agent = _make_task_agent(output_schema=output_schema)
+    tool = FinishTaskTool(task_agent=agent)
+    declaration = tool._get_declaration()
+
+    assert declaration is not None
+    assert declaration.name == FINISH_TASK_TOOL_NAME
+    schema = declaration.parameters_json_schema
+    assert schema is not None
+    assert 'result' in schema.get('properties', {})
+
+  @pytest.mark.parametrize(
+      'output_schema, args, expected_output',
+      [
+          (
+              SampleOutputSchema,
+              {'result': 'done', 'count': 5},
+              {'result': 'done', 'count': 5},
+          ),
+          (dict[str, Any], {'key': 'value'}, {'key': 'value'}),
+          (str, {'result': 'hello'}, 'hello'),
+          (int, {'result': 42}, 42),
+          (bool, {'result': True}, True),
+          (float, {'result': 3.14}, 3.14),
+          (list[str], {'result': ['a', 'b', 'c']}, ['a', 'b', 'c']),
+          (list[int], {'result': [1, 2, 3]}, [1, 2, 3]),
+          (
+              list[SampleOutputSchema],
+              {'result': [{'result': 'ok', 'count': 1}]},
+              [{'result': 'ok', 'count': 1}],
+          ),
+          (list[str], {'result': []}, []),
+      ],
+      ids=[
+          'BaseModel',
+          'dict',
+          'str',
+          'int',
+          'bool',
+          'float',
+          'list_str',
+          'list_int',
+          'list_BaseModel',
+          'list_str_empty',
+      ],
+  )
+  @pytest.mark.asyncio
+  async def test_run_async(self, output_schema, args, expected_output):
+    """Verify run_async validates and extracts output for each schema type."""
+    agent = _make_task_agent(output_schema=output_schema)
+    tool = FinishTaskTool(task_agent=agent)
+    mock_tool_context = mock.MagicMock()
+
+    result = await tool.run_async(
+        args=args,
+        tool_context=mock_tool_context,
+    )
+
+    # Task Delegation API: tool no longer writes actions.finish_task. The
+    # LlmAgent task wrapper sniffs the FC's `output` arg directly to
+    # set event.output. The tool just performs schema validation.
+    assert result == 'Task completed.'
+    del expected_output  # validated_output is no longer surfaced via actions
+
+  def test_get_declaration_list_basemodel_defs_at_root(self):
+    """Non-object schemas with $defs should hoist $defs to root level."""
+    agent = _make_task_agent(output_schema=list[SampleOutputSchema])
+    tool = FinishTaskTool(task_agent=agent)
+    declaration = tool._get_declaration()
+
+    schema = declaration.parameters_json_schema
+    # $defs must be at the root, not nested inside properties.result
+    assert '$defs' in schema
+    assert 'SampleOutputSchema' in schema['$defs']
+    # The nested result schema should not contain $defs
+    assert '$defs' not in schema['properties']['result']

tests/unittests/agents/test_agent_clone.py

@@ -0,0 +1,573 @@
+# 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.
+
+"""Testings for the clone functionality of agents."""
+
+from typing import Any
+from typing import cast
+from typing import Iterable
+
+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
+import pytest
+
+
+def test_llm_agent_clone():
+  """Test cloning an LLM agent."""
+  # Create an LLM agent
+  original = LlmAgent(
+      name="llm_agent",
+      description="An LLM agent",
+      instruction="You are a helpful assistant.",
+  )
+
+  # Clone it with name update
+  cloned = original.clone(update={"name": "cloned_llm_agent"})
+
+  # Verify the clone
+  assert cloned.name == "cloned_llm_agent"
+  assert cloned.description == "An LLM agent"
+  assert cloned.instruction == "You are a helpful assistant."
+  assert cloned.parent_agent is None
+  assert len(cloned.sub_agents) == 0
+  assert isinstance(cloned, LlmAgent)
+
+  # Verify the original is unchanged
+  assert original.name == "llm_agent"
+  assert original.instruction == "You are a helpful assistant."
+
+
+def test_agent_with_sub_agents():
+  """Test cloning an agent that has sub-agents."""
+  # Create sub-agents
+  sub_agent1 = LlmAgent(name="sub_agent1", description="First sub-agent")
+  sub_agent2 = LlmAgent(name="sub_agent2", description="Second sub-agent")
+
+  # Create a parent agent with sub-agents
+  original = SequentialAgent(
+      name="parent_agent",
+      description="Parent agent with sub-agents",
+      sub_agents=[sub_agent1, sub_agent2],
+  )
+
+  # Clone it with name update
+  cloned = original.clone(update={"name": "cloned_parent"})
+
+  # Verify the clone has sub-agents (deep copy behavior)
+  assert cloned.name == "cloned_parent"
+  assert cloned.description == "Parent agent with sub-agents"
+  assert cloned.parent_agent is None
+  assert len(cloned.sub_agents) == 2
+
+  # Sub-agents should be cloned with their original names
+  assert cloned.sub_agents[0].name == "sub_agent1"
+  assert cloned.sub_agents[1].name == "sub_agent2"
+
+  # Sub-agents should have the cloned agent as their parent
+  assert cloned.sub_agents[0].parent_agent == cloned
+  assert cloned.sub_agents[1].parent_agent == cloned
+
+  # Sub-agents should be different objects from the original
+  assert cloned.sub_agents[0] is not original.sub_agents[0]
+  assert cloned.sub_agents[1] is not original.sub_agents[1]
+
+  # Verify the original still has sub-agents
+  assert original.name == "parent_agent"
+  assert len(original.sub_agents) == 2
+  assert original.sub_agents[0].name == "sub_agent1"
+  assert original.sub_agents[1].name == "sub_agent2"
+  assert original.sub_agents[0].parent_agent == original
+  assert original.sub_agents[1].parent_agent == original
+
+
+def test_three_level_nested_agent():
+  """Test cloning a three-level nested agent to verify recursive cloning logic."""
+  # Create third-level agents (leaf nodes)
+  leaf_agent1 = LlmAgent(name="leaf1", description="First leaf agent")
+  leaf_agent2 = LlmAgent(name="leaf2", description="Second leaf agent")
+
+  # Create second-level agents
+  middle_agent1 = SequentialAgent(
+      name="middle1", description="First middle agent", sub_agents=[leaf_agent1]
+  )
+  middle_agent2 = ParallelAgent(
+      name="middle2",
+      description="Second middle agent",
+      sub_agents=[leaf_agent2],
+  )
+
+  # Create top-level agent
+  root_agent = LoopAgent(
+      name="root_agent",
+      description="Root agent with three levels",
+      max_iterations=5,
+      sub_agents=[middle_agent1, middle_agent2],
+  )
+
+  # Clone the root agent
+  cloned_root = root_agent.clone(update={"name": "cloned_root"})
+
+  # Verify root level
+  assert cloned_root.name == "cloned_root"
+  assert cloned_root.description == "Root agent with three levels"
+  assert cloned_root.max_iterations == 5
+  assert cloned_root.parent_agent is None
+  assert len(cloned_root.sub_agents) == 2
+  assert isinstance(cloned_root, LoopAgent)
+
+  # Verify middle level
+  cloned_middle1 = cloned_root.sub_agents[0]
+  cloned_middle2 = cloned_root.sub_agents[1]
+
+  assert cloned_middle1.name == "middle1"
+  assert cloned_middle1.description == "First middle agent"
+  assert cloned_middle1.parent_agent == cloned_root
+  assert len(cloned_middle1.sub_agents) == 1
+  assert isinstance(cloned_middle1, SequentialAgent)
+
+  assert cloned_middle2.name == "middle2"
+  assert cloned_middle2.description == "Second middle agent"
+  assert cloned_middle2.parent_agent == cloned_root
+  assert len(cloned_middle2.sub_agents) == 1
+  assert isinstance(cloned_middle2, ParallelAgent)
+
+  # Verify leaf level
+  cloned_leaf1 = cloned_middle1.sub_agents[0]
+  cloned_leaf2 = cloned_middle2.sub_agents[0]
+
+  assert cloned_leaf1.name == "leaf1"
+  assert cloned_leaf1.description == "First leaf agent"
+  assert cloned_leaf1.parent_agent == cloned_middle1
+  assert len(cloned_leaf1.sub_agents) == 0
+  assert isinstance(cloned_leaf1, LlmAgent)
+
+  assert cloned_leaf2.name == "leaf2"
+  assert cloned_leaf2.description == "Second leaf agent"
+  assert cloned_leaf2.parent_agent == cloned_middle2
+  assert len(cloned_leaf2.sub_agents) == 0
+  assert isinstance(cloned_leaf2, LlmAgent)
+
+  # Verify all objects are different from originals
+  assert cloned_root is not root_agent
+  assert cloned_middle1 is not middle_agent1
+  assert cloned_middle2 is not middle_agent2
+  assert cloned_leaf1 is not leaf_agent1
+  assert cloned_leaf2 is not leaf_agent2
+
+  # Verify original structure is unchanged
+  assert root_agent.name == "root_agent"
+  assert root_agent.sub_agents[0].name == "middle1"
+  assert root_agent.sub_agents[1].name == "middle2"
+  assert root_agent.sub_agents[0].sub_agents[0].name == "leaf1"
+  assert root_agent.sub_agents[1].sub_agents[0].name == "leaf2"
+
+
+def test_multiple_clones():
+  """Test creating multiple clones with automatic naming."""
+  # Create multiple agents and clone each one
+  original = LlmAgent(
+      name="original_agent", description="Agent for multiple cloning"
+  )
+
+  # Test multiple clones from the same original
+  clone1 = original.clone(update={"name": "clone1"})
+  clone2 = original.clone(update={"name": "clone2"})
+
+  assert clone1.name == "clone1"
+  assert clone2.name == "clone2"
+  assert clone1 is not clone2
+
+
+def test_clone_with_complex_configuration():
+  """Test cloning an agent with complex configuration."""
+  # Create an LLM agent with various configurations
+  original = LlmAgent(
+      name="complex_agent",
+      description="A complex agent with many settings",
+      instruction="You are a specialized assistant.",
+      global_instruction="Always be helpful and accurate.",
+      disallow_transfer_to_parent=True,
+      disallow_transfer_to_peers=True,
+      include_contents="none",
+  )
+
+  # Clone it with name update
+  cloned = original.clone(update={"name": "complex_clone"})
+
+  # Verify all configurations are preserved
+  assert cloned.name == "complex_clone"
+  assert cloned.description == "A complex agent with many settings"
+  assert cloned.instruction == "You are a specialized assistant."
+  assert cloned.global_instruction == "Always be helpful and accurate."
+  assert cloned.disallow_transfer_to_parent is True
+  assert cloned.disallow_transfer_to_peers is True
+  assert cloned.include_contents == "none"
+
+  # Verify parent and sub-agents are set
+  assert cloned.parent_agent is None
+  assert len(cloned.sub_agents) == 0
+
+
+def test_clone_without_updates():
+  """Test cloning without providing updates (should use original values)."""
+  original = LlmAgent(name="test_agent", description="Test agent")
+
+  cloned = original.clone()
+
+  assert cloned.name == "test_agent"
+  assert cloned.description == "Test agent"
+
+
+def test_clone_with_multiple_updates():
+  """Test cloning with multiple field updates."""
+  original = LlmAgent(
+      name="original_agent",
+      description="Original description",
+      instruction="Original instruction",
+  )
+
+  cloned = original.clone(
+      update={
+          "name": "updated_agent",
+          "description": "Updated description",
+          "instruction": "Updated instruction",
+      }
+  )
+
+  assert cloned.name == "updated_agent"
+  assert cloned.description == "Updated description"
+  assert cloned.instruction == "Updated instruction"
+
+
+def test_clone_with_sub_agents_deep_copy():
+  """Test cloning with deep copy of sub-agents."""
+  # Create an agent with sub-agents
+  sub_agent = LlmAgent(name="sub_agent", description="Sub agent")
+  original = LlmAgent(
+      name="root_agent",
+      description="Root agent",
+      sub_agents=[sub_agent],
+  )
+
+  # Clone with deep copy
+  cloned = original.clone(update={"name": "cloned_root_agent"})
+  assert cloned.name == "cloned_root_agent"
+  assert cloned.sub_agents[0].name == "sub_agent"
+  assert cloned.sub_agents[0].parent_agent == cloned
+  assert cloned.sub_agents[0] is not original.sub_agents[0]
+
+
+def test_clone_invalid_field():
+  """Test that cloning with invalid fields raises an error."""
+  original = LlmAgent(name="test_agent", description="Test agent")
+
+  with pytest.raises(ValueError, match="Cannot update nonexistent fields"):
+    original.clone(update={"invalid_field": "value"})
+
+
+def test_clone_parent_agent_field():
+  """Test that cloning with parent_agent field raises an error."""
+  original = LlmAgent(name="test_agent", description="Test agent")
+
+  with pytest.raises(
+      ValueError, match="Cannot update `parent_agent` field in clone"
+  ):
+    original.clone(update={"parent_agent": None})
+
+
+def test_clone_preserves_agent_type():
+  """Test that cloning preserves the specific agent type."""
+  # Test LlmAgent
+  llm_original = LlmAgent(name="llm_test")
+  llm_cloned = llm_original.clone()
+  assert isinstance(llm_cloned, LlmAgent)
+
+  # Test SequentialAgent
+  seq_original = SequentialAgent(
+      name="seq_test", sub_agents=[LlmAgent(name="dummy_seq")]
+  )
+  seq_cloned = seq_original.clone()
+  assert isinstance(seq_cloned, SequentialAgent)
+
+  # Test ParallelAgent
+  par_original = ParallelAgent(
+      name="par_test", sub_agents=[LlmAgent(name="dummy_par")]
+  )
+  par_cloned = par_original.clone()
+  assert isinstance(par_cloned, ParallelAgent)
+
+  # Test LoopAgent
+  loop_original = LoopAgent(
+      name="loop_test", sub_agents=[LlmAgent(name="dummy_loop")]
+  )
+  loop_cloned = loop_original.clone()
+  assert isinstance(loop_cloned, LoopAgent)
+
+
+def test_clone_with_agent_specific_fields():
+  # Test LoopAgent
+  dummy = LlmAgent(name="dummy")
+  loop_original = LoopAgent(name="loop_test", sub_agents=[dummy])
+  loop_cloned = loop_original.clone({"max_iterations": 10})
+  assert isinstance(loop_cloned, LoopAgent)
+  assert loop_cloned.max_iterations == 10
+
+
+def test_clone_with_none_update():
+  """Test cloning with explicit None update parameter."""
+  original = LlmAgent(name="test_agent", description="Test agent")
+
+  cloned = original.clone(update=None)
+
+  assert cloned.name == "test_agent"
+  assert cloned.description == "Test agent"
+  assert cloned is not original
+
+
+def test_clone_with_empty_update():
+  """Test cloning with empty update dictionary."""
+  original = LlmAgent(name="test_agent", description="Test agent")
+
+  cloned = original.clone(update={})
+
+  assert cloned.name == "test_agent"
+  assert cloned.description == "Test agent"
+  assert cloned is not original
+
+
+def test_clone_with_sub_agents_update():
+  """Test cloning with sub_agents provided in update."""
+  # Create original sub-agents
+  original_sub1 = LlmAgent(name="original_sub1", description="Original sub 1")
+  original_sub2 = LlmAgent(name="original_sub2", description="Original sub 2")
+
+  # Create new sub-agents for the update
+  new_sub1 = LlmAgent(name="new_sub1", description="New sub 1")
+  new_sub2 = LlmAgent(name="new_sub2", description="New sub 2")
+
+  # Create original agent with sub-agents
+  original = SequentialAgent(
+      name="original_agent",
+      description="Original agent",
+      sub_agents=[original_sub1, original_sub2],
+  )
+
+  # Clone with sub_agents update
+  cloned = original.clone(
+      update={"name": "cloned_agent", "sub_agents": [new_sub1, new_sub2]}
+  )
+
+  # Verify the clone uses the new sub-agents
+  assert cloned.name == "cloned_agent"
+  assert len(cloned.sub_agents) == 2
+  assert cloned.sub_agents[0].name == "new_sub1"
+  assert cloned.sub_agents[1].name == "new_sub2"
+  assert cloned.sub_agents[0].parent_agent == cloned
+  assert cloned.sub_agents[1].parent_agent == cloned
+
+  # Verify original is unchanged
+  assert original.name == "original_agent"
+  assert len(original.sub_agents) == 2
+  assert original.sub_agents[0].name == "original_sub1"
+  assert original.sub_agents[1].name == "original_sub2"
+
+
+def _check_lists_contain_same_contents(*lists: Iterable[list[Any]]) -> None:
+  """Assert that all provided lists contain the same elements."""
+  if lists:
+    first_list = lists[0]
+    assert all(len(lst) == len(first_list) for lst in lists)
+    for idx, elem in enumerate(first_list):
+      assert all(lst[idx] is elem for lst in lists)
+
+
+def test_clone_shallow_copies_lists():
+  """Test that cloning shallow copies fields stored as lists."""
+  # Define the list fields
+  before_agent_callback = [lambda *args, **kwargs: None]
+  after_agent_callback = [lambda *args, **kwargs: None]
+  before_model_callback = [lambda *args, **kwargs: None]
+  after_model_callback = [lambda *args, **kwargs: None]
+  before_tool_callback = [lambda *args, **kwargs: None]
+  after_tool_callback = [lambda *args, **kwargs: None]
+  tools = [lambda *args, **kwargs: None]
+
+  # Create the original agent with list fields
+  original = LlmAgent(
+      name="original_agent",
+      description="Original agent",
+      before_agent_callback=before_agent_callback,
+      after_agent_callback=after_agent_callback,
+      before_model_callback=before_model_callback,
+      after_model_callback=after_model_callback,
+      before_tool_callback=before_tool_callback,
+      after_tool_callback=after_tool_callback,
+      tools=tools,
+  )
+
+  # Clone the agent
+  cloned = original.clone()
+
+  # Verify the lists are copied
+  assert original.before_agent_callback is not cloned.before_agent_callback
+  assert original.after_agent_callback is not cloned.after_agent_callback
+  assert original.before_model_callback is not cloned.before_model_callback
+  assert original.after_model_callback is not cloned.after_model_callback
+  assert original.before_tool_callback is not cloned.before_tool_callback
+  assert original.after_tool_callback is not cloned.after_tool_callback
+  assert original.tools is not cloned.tools
+
+  # Verify the list copies are shallow
+  _check_lists_contain_same_contents(
+      before_agent_callback,
+      original.before_agent_callback,
+      cloned.before_agent_callback,
+  )
+  _check_lists_contain_same_contents(
+      after_agent_callback,
+      original.after_agent_callback,
+      cloned.after_agent_callback,
+  )
+  _check_lists_contain_same_contents(
+      before_model_callback,
+      original.before_model_callback,
+      cloned.before_model_callback,
+  )
+  _check_lists_contain_same_contents(
+      after_model_callback,
+      original.after_model_callback,
+      cloned.after_model_callback,
+  )
+  _check_lists_contain_same_contents(
+      before_tool_callback,
+      original.before_tool_callback,
+      cloned.before_tool_callback,
+  )
+  _check_lists_contain_same_contents(
+      after_tool_callback,
+      original.after_tool_callback,
+      cloned.after_tool_callback,
+  )
+  _check_lists_contain_same_contents(tools, original.tools, cloned.tools)
+
+
+def test_clone_shallow_copies_lists_with_sub_agents():
+  """Test that cloning recursively shallow copies fields stored as lists."""
+  # Define the list fields for the sub-agent
+  before_agent_callback = [lambda *args, **kwargs: None]
+  after_agent_callback = [lambda *args, **kwargs: None]
+  before_model_callback = [lambda *args, **kwargs: None]
+  after_model_callback = [lambda *args, **kwargs: None]
+  before_tool_callback = [lambda *args, **kwargs: None]
+  after_tool_callback = [lambda *args, **kwargs: None]
+  tools = [lambda *args, **kwargs: None]
+
+  # Create the original sub-agent with list fields and the top-level agent
+  sub_agents = [
+      LlmAgent(
+          name="sub_agent",
+          description="Sub agent",
+          before_agent_callback=before_agent_callback,
+          after_agent_callback=after_agent_callback,
+          before_model_callback=before_model_callback,
+          after_model_callback=after_model_callback,
+          before_tool_callback=before_tool_callback,
+          after_tool_callback=after_tool_callback,
+          tools=tools,
+      )
+  ]
+  original = LlmAgent(
+      name="original_agent",
+      description="Original agent",
+      sub_agents=sub_agents,
+  )
+
+  # Clone the top-level agent
+  cloned = original.clone()
+
+  # Verify the sub_agents list is copied for the top-level agent
+  assert original.sub_agents is not cloned.sub_agents
+
+  # Retrieve the sub-agent for the original and cloned top-level agent
+  original_sub_agent = cast(LlmAgent, original.sub_agents[0])
+  cloned_sub_agent = cast(LlmAgent, cloned.sub_agents[0])
+
+  # Verify the lists are copied for the sub-agent
+  assert (
+      original_sub_agent.before_agent_callback
+      is not cloned_sub_agent.before_agent_callback
+  )
+  assert (
+      original_sub_agent.after_agent_callback
+      is not cloned_sub_agent.after_agent_callback
+  )
+  assert (
+      original_sub_agent.before_model_callback
+      is not cloned_sub_agent.before_model_callback
+  )
+  assert (
+      original_sub_agent.after_model_callback
+      is not cloned_sub_agent.after_model_callback
+  )
+  assert (
+      original_sub_agent.before_tool_callback
+      is not cloned_sub_agent.before_tool_callback
+  )
+  assert (
+      original_sub_agent.after_tool_callback
+      is not cloned_sub_agent.after_tool_callback
+  )
+  assert original_sub_agent.tools is not cloned_sub_agent.tools
+
+  # Verify the list copies are shallow for the sub-agent
+  _check_lists_contain_same_contents(
+      before_agent_callback,
+      original_sub_agent.before_agent_callback,
+      cloned_sub_agent.before_agent_callback,
+  )
+  _check_lists_contain_same_contents(
+      after_agent_callback,
+      original_sub_agent.after_agent_callback,
+      cloned_sub_agent.after_agent_callback,
+  )
+  _check_lists_contain_same_contents(
+      before_model_callback,
+      original_sub_agent.before_model_callback,
+      cloned_sub_agent.before_model_callback,
+  )
+  _check_lists_contain_same_contents(
+      after_model_callback,
+      original_sub_agent.after_model_callback,
+      cloned_sub_agent.after_model_callback,
+  )
+  _check_lists_contain_same_contents(
+      before_tool_callback,
+      original_sub_agent.before_tool_callback,
+      cloned_sub_agent.before_tool_callback,
+  )
+  _check_lists_contain_same_contents(
+      after_tool_callback,
+      original_sub_agent.after_tool_callback,
+      cloned_sub_agent.after_tool_callback,
+  )
+  _check_lists_contain_same_contents(
+      tools, original_sub_agent.tools, cloned_sub_agent.tools
+  )
+
+
+if __name__ == "__main__":
+  # Run a specific test for debugging
+  test_three_level_nested_agent()

tests/unittests/agents/test_agent_config.py

@@ -0,0 +1,618 @@
+# 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.
+
+import ntpath
+import os
+from pathlib import Path
+from textwrap import dedent
+from typing import Literal
+from typing import Type
+from unittest import mock
+
+from google.adk.agents import config_agent_utils
+from google.adk.agents.agent_config import AgentConfig
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.base_agent_config import BaseAgentConfig
+from google.adk.agents.common_configs import AgentRefConfig
+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.models.lite_llm import LiteLlm
+import pytest
+import yaml
+
+
+def test_agent_config_discriminator_default_is_llm_agent(tmp_path: Path):
+  yaml_content = """\
+name: search_agent
+model: gemini-2.5-flash
+description: a sample description
+instruction: a fake instruction
+tools:
+  - name: google_search
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, LlmAgent)
+  assert config.root.agent_class == "LlmAgent"
+
+
+@pytest.mark.parametrize(
+    "agent_class_value",
+    [
+        "LlmAgent",
+        "google.adk.agents.LlmAgent",
+        "google.adk.agents.llm_agent.LlmAgent",
+    ],
+)
+def test_agent_config_discriminator_llm_agent(
+    agent_class_value: str, tmp_path: Path
+):
+  yaml_content = f"""\
+agent_class: {agent_class_value}
+name: search_agent
+model: gemini-2.5-flash
+description: a sample description
+instruction: a fake instruction
+tools:
+  - name: google_search
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, LlmAgent)
+  assert config.root.agent_class == agent_class_value
+
+
+@pytest.mark.parametrize(
+    "agent_class_value",
+    [
+        "LoopAgent",
+        "google.adk.agents.LoopAgent",
+        "google.adk.agents.loop_agent.LoopAgent",
+    ],
+)
+def test_agent_config_discriminator_loop_agent(
+    agent_class_value: str, tmp_path: Path
+):
+  yaml_content = f"""\
+agent_class: {agent_class_value}
+name: CodePipelineAgent
+description: Executes a sequence of code writing, reviewing, and refactoring.
+sub_agents: []
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, LoopAgent)
+  assert config.root.agent_class == agent_class_value
+
+
+@pytest.mark.parametrize(
+    "agent_class_value",
+    [
+        "ParallelAgent",
+        "google.adk.agents.ParallelAgent",
+        "google.adk.agents.parallel_agent.ParallelAgent",
+    ],
+)
+def test_agent_config_discriminator_parallel_agent(
+    agent_class_value: str, tmp_path: Path
+):
+  yaml_content = f"""\
+agent_class: {agent_class_value}
+name: CodePipelineAgent
+description: Executes a sequence of code writing, reviewing, and refactoring.
+sub_agents: []
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, ParallelAgent)
+  assert config.root.agent_class == agent_class_value
+
+
+@pytest.mark.parametrize(
+    "agent_class_value",
+    [
+        "SequentialAgent",
+        "google.adk.agents.SequentialAgent",
+        "google.adk.agents.sequential_agent.SequentialAgent",
+    ],
+)
+def test_agent_config_discriminator_sequential_agent(
+    agent_class_value: str, tmp_path: Path
+):
+  yaml_content = f"""\
+agent_class: {agent_class_value}
+name: CodePipelineAgent
+description: Executes a sequence of code writing, reviewing, and refactoring.
+sub_agents: []
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, SequentialAgent)
+  assert config.root.agent_class == agent_class_value
+
+
+@pytest.mark.parametrize(
+    ("agent_class_value", "expected_agent_type"),
+    [
+        ("LoopAgent", LoopAgent),
+        ("google.adk.agents.LoopAgent", LoopAgent),
+        ("google.adk.agents.loop_agent.LoopAgent", LoopAgent),
+        ("ParallelAgent", ParallelAgent),
+        ("google.adk.agents.ParallelAgent", ParallelAgent),
+        ("google.adk.agents.parallel_agent.ParallelAgent", ParallelAgent),
+        ("SequentialAgent", SequentialAgent),
+        ("google.adk.agents.SequentialAgent", SequentialAgent),
+        ("google.adk.agents.sequential_agent.SequentialAgent", SequentialAgent),
+    ],
+)
+def test_agent_config_discriminator_with_sub_agents(
+    agent_class_value: str, expected_agent_type: Type[BaseAgent], tmp_path: Path
+):
+  # Create sub-agent config files
+  sub_agent_dir = tmp_path / "sub_agents"
+  sub_agent_dir.mkdir()
+  sub_agent_config = """\
+name: sub_agent_{index}
+model: gemini-2.5-flash
+description: a sub agent
+instruction: sub agent instruction
+"""
+  (sub_agent_dir / "sub_agent1.yaml").write_text(
+      sub_agent_config.format(index=1)
+  )
+  (sub_agent_dir / "sub_agent2.yaml").write_text(
+      sub_agent_config.format(index=2)
+  )
+  yaml_content = f"""\
+agent_class: {agent_class_value}
+name: main_agent
+description: main agent with sub agents
+sub_agents:
+  - config_path: sub_agents/sub_agent1.yaml
+  - config_path: sub_agents/sub_agent2.yaml
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, expected_agent_type)
+  assert config.root.agent_class == agent_class_value
+
+
+@pytest.mark.parametrize(
+    ("agent_class_value", "expected_agent_type"),
+    [
+        ("LlmAgent", LlmAgent),
+        ("google.adk.agents.LlmAgent", LlmAgent),
+        ("google.adk.agents.llm_agent.LlmAgent", LlmAgent),
+    ],
+)
+def test_agent_config_discriminator_llm_agent_with_sub_agents(
+    agent_class_value: str, expected_agent_type: Type[BaseAgent], tmp_path: Path
+):
+  # Create sub-agent config files
+  sub_agent_dir = tmp_path / "sub_agents"
+  sub_agent_dir.mkdir()
+  sub_agent_config = """\
+name: sub_agent_{index}
+model: gemini-2.5-flash
+description: a sub agent
+instruction: sub agent instruction
+"""
+  (sub_agent_dir / "sub_agent1.yaml").write_text(
+      sub_agent_config.format(index=1)
+  )
+  (sub_agent_dir / "sub_agent2.yaml").write_text(
+      sub_agent_config.format(index=2)
+  )
+  yaml_content = f"""\
+agent_class: {agent_class_value}
+name: main_agent
+model: gemini-2.5-flash
+description: main agent with sub agents
+instruction: main agent instruction
+sub_agents:
+  - config_path: sub_agents/sub_agent1.yaml
+  - config_path: sub_agents/sub_agent2.yaml
+"""
+  config_file = tmp_path / "test_config.yaml"
+  config_file.write_text(yaml_content)
+
+  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
+  agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, expected_agent_type)
+  assert config.root.agent_class == agent_class_value
+
+
+def test_agent_config_model_code_resolves_preconfigured_client(tmp_path: Path):
+  """model_code references a pre-built model instance by fully qualified name.
+
+  Configured clients (custom api_base, etc.) must be constructed in Python
+  and referenced from YAML; YAML cannot pass constructor arguments.
+  """
+  preconfigured = LiteLlm(
+      model="kimi/k2", api_base="https://proxy.litellm.ai/v1"
+  )
+
+  yaml_content = """\
+name: managed_api_agent
+description: Agent using LiteLLM managed endpoint
+instruction: Respond concisely.
+model_code:
+  name: my_library.clients.my_litellm
+"""
+  config_file = tmp_path / "litellm_agent.yaml"
+  config_file.write_text(yaml_content)
+
+  with mock.patch.object(
+      config_agent_utils,
+      "resolve_code_reference",
+      return_value=preconfigured,
+  ):
+    agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, LlmAgent)
+  assert agent.model is preconfigured
+
+
+def test_agent_config_discriminator_custom_agent():
+  class MyCustomAgentConfig(BaseAgentConfig):
+    agent_class: Literal["mylib.agents.MyCustomAgent"] = (
+        "mylib.agents.MyCustomAgent"
+    )
+    other_field: str
+
+  yaml_content = """\
+agent_class: mylib.agents.MyCustomAgent
+name: CodePipelineAgent
+description: Executes a sequence of code writing, reviewing, and refactoring.
+other_field: other value
+"""
+  config_data = yaml.safe_load(yaml_content)
+
+  config = AgentConfig.model_validate(config_data)
+
+  # pylint: disable=unidiomatic-typecheck Needs exact class matching.
+  assert type(config.root) is BaseAgentConfig
+  assert config.root.agent_class == "mylib.agents.MyCustomAgent"
+  assert config.root.model_extra == {"other_field": "other value"}
+
+  my_custom_config = MyCustomAgentConfig.model_validate(
+      config.root.model_dump()
+  )
+  assert my_custom_config.other_field == "other value"
+
+
+def test_from_config_passes_extra_yaml_fields_to_custom_agent_constructor(
+    tmp_path: Path,
+):
+  """Custom agent fields in YAML reach the constructor without a custom config_type.
+
+  Mirrors the 1.x AgentConfigMapper behavior: a custom agent subclass with
+  extra Pydantic fields declared on the agent (not on a config_type) can
+  populate those fields directly from YAML.
+  """
+
+  class MyCustomAgent(BaseAgent):
+    custom_field: str = ""
+
+  yaml_content = """\
+agent_class: mylib.agents.MyCustomAgent
+name: custom_agent
+description: a custom agent
+custom_field: hello from yaml
+"""
+  config_file = tmp_path / "custom_agent.yaml"
+  config_file.write_text(yaml_content)
+
+  with mock.patch.object(
+      config_agent_utils,
+      "resolve_fully_qualified_name",
+      return_value=MyCustomAgent,
+  ):
+    agent = config_agent_utils.from_config(str(config_file))
+
+  assert isinstance(agent, MyCustomAgent)
+  assert agent.custom_field == "hello from yaml"
+
+
+def test_from_config_ignores_extra_yaml_fields_not_on_agent(tmp_path: Path):
+  """Extra YAML keys that don't map to constructor params are silently dropped."""
+
+  class MyCustomAgent(BaseAgent):
+    custom_field: str = ""
+
+  yaml_content = """\
+agent_class: mylib.agents.MyCustomAgent
+name: custom_agent
+description: a custom agent
+custom_field: kept
+unknown_field: dropped
+"""
+  config_file = tmp_path / "custom_agent.yaml"
+  config_file.write_text(yaml_content)
+
+  with mock.patch.object(
+      config_agent_utils,
+      "resolve_fully_qualified_name",
+      return_value=MyCustomAgent,
+  ):
+    agent = config_agent_utils.from_config(str(config_file))
+
+  assert agent.custom_field == "kept"
+  assert not hasattr(agent, "unknown_field")
+
+
+@pytest.mark.parametrize(
+    ("config_rel_path", "child_rel_path", "child_name", "instruction"),
+    [
+        (
+            Path("main.yaml"),
+            Path("sub_agents/child.yaml"),
+            "child_agent",
+            "I am a child agent",
+        ),
+        (
+            Path("level1/level2/nested_main.yaml"),
+            Path("sub/nested_child.yaml"),
+            "nested_child",
+            "I am nested",
+        ),
+    ],
+)
+def test_resolve_agent_reference_resolves_relative_paths(
+    config_rel_path: Path,
+    child_rel_path: Path,
+    child_name: str,
+    instruction: str,
+    tmp_path: Path,
+):
+  """Verify resolve_agent_reference resolves relative sub-agent paths."""
+  config_file = tmp_path / config_rel_path
+  config_file.parent.mkdir(parents=True, exist_ok=True)
+
+  child_config_path = config_file.parent / child_rel_path
+  child_config_path.parent.mkdir(parents=True, exist_ok=True)
+  child_config_path.write_text(dedent(f"""
+          agent_class: LlmAgent
+          name: {child_name}
+          model: gemini-2.5-flash
+          instruction: {instruction}
+          """).lstrip())
+
+  config_file.write_text(dedent(f"""
+          agent_class: LlmAgent
+          name: main_agent
+          model: gemini-2.5-flash
+          instruction: I am the main agent
+          sub_agents:
+            - config_path: {child_rel_path.as_posix()}
+          """).lstrip())
+
+  ref_config = AgentRefConfig(config_path=child_rel_path.as_posix())
+  agent = config_agent_utils.resolve_agent_reference(
+      ref_config, str(config_file)
+  )
+
+  assert agent.name == child_name
+
+  config_dir = os.path.dirname(str(config_file.resolve()))
+  assert config_dir == str(config_file.parent.resolve())
+
+  expected_child_path = os.path.join(config_dir, *child_rel_path.parts)
+  assert os.path.exists(expected_child_path)
+
+
+def test_resolve_agent_reference_uses_windows_dirname():
+  """Ensure Windows-style config references resolve via os.path.dirname."""
+  ref_config = AgentRefConfig(config_path="sub\\child.yaml")
+  recorded: dict[str, str] = {}
+
+  def fake_from_config(path: str):
+    recorded["path"] = path
+    return "sentinel"
+
+  with (
+      mock.patch.object(
+          config_agent_utils,
+          "from_config",
+          autospec=True,
+          side_effect=fake_from_config,
+      ),
+      mock.patch.object(config_agent_utils.os, "path", ntpath),
+  ):
+    referencing = r"C:\workspace\agents\main.yaml"
+    result = config_agent_utils.resolve_agent_reference(ref_config, referencing)
+
+  expected_path = ntpath.join(
+      ntpath.dirname(referencing), ref_config.config_path
+  )
+  assert result == "sentinel"
+  assert recorded["path"] == expected_path
+
+
+def test_resolve_agent_reference_blocks_absolute_path():
+  """Verify resolve_agent_reference raises ValueError for absolute paths."""
+  ref_config = AgentRefConfig(config_path="/etc/passwd")
+  with pytest.raises(
+      ValueError,
+      match="Absolute paths are not allowed in AgentRefConfig config_path",
+  ):
+    config_agent_utils.resolve_agent_reference(
+        ref_config, "/workspace/main.yaml"
+    )
+
+
+def test_resolve_agent_reference_blocks_path_traversal():
+  """Verify resolve_agent_reference raises ValueError for path traversal."""
+  ref_config = AgentRefConfig(config_path="../outside.yaml")
+  with pytest.raises(ValueError, match="Path traversal detected"):
+    config_agent_utils.resolve_agent_reference(
+        ref_config, "/workspace/agents/main.yaml"
+    )
+
+
+# --- Security tests: module blocklist for YAML agent config code references ---
+
+
+def test_resolve_code_reference_blocks_os_when_enforced():
+  """Verify resolve_code_reference blocks os module directly."""
+  from google.adk.agents.common_configs import CodeConfig
+
+  with pytest.raises(ValueError, match="Blocked module reference"):
+    config_agent_utils.resolve_code_reference(CodeConfig(name="os.system"))
+
+
+def test_resolve_fully_qualified_name_blocks_subprocess_when_enforced():
+  """Verify resolve_fully_qualified_name blocks subprocess module.
+
+  resolve_fully_qualified_name wraps all exceptions in
+  ValueError("Invalid fully qualified name: ..."), so we check the wrapper
+  and verify the __cause__ carries the blocklist message.
+  """
+  with pytest.raises(
+      ValueError, match="Invalid fully qualified name"
+  ) as exc_info:
+    config_agent_utils.resolve_fully_qualified_name("subprocess.Popen")
+  assert "Blocked module reference" in str(exc_info.value.__cause__)
+
+
+def test_allowed_module_passes_when_enforced(tmp_path: Path):
+  """Verify that google.adk modules are NOT blocked by the module denylist."""
+  # This should NOT raise — google.adk modules must remain allowed
+  result = config_agent_utils.resolve_fully_qualified_name(
+      "google.adk.agents.llm_agent.LlmAgent"
+  )
+  assert result is LlmAgent
+
+
+@pytest.mark.parametrize(
+    "blocked_module",
+    [
+        "os.system",
+        "subprocess.call",
+        "builtins.exec",
+    ],
+)
+def test_resolve_agent_code_reference_blocks_when_enforced(
+    blocked_module: str,
+):
+  """Verify _resolve_agent_code_reference blocks dangerous modules."""
+  with pytest.raises(ValueError, match="Blocked module reference"):
+    config_agent_utils._resolve_agent_code_reference(blocked_module)
+
+
+@pytest.mark.parametrize(
+    "blocked_ref",
+    [
+        "os.system",
+        "subprocess.call",
+        "builtins.exec",
+        "pickle.loads",
+    ],
+)
+def test_resolve_tools_blocks_dangerous_modules(blocked_ref: str):
+  """Verify _resolve_tools blocks dangerous modules for user-defined tools."""
+  from google.adk.agents.llm_agent import LlmAgent
+  from google.adk.tools.tool_configs import ToolConfig
+
+  tool_config = ToolConfig(name=blocked_ref)
+  with pytest.raises(ValueError, match="Blocked module reference"):
+    LlmAgent._resolve_tools([tool_config], "/fake/path.yaml")
+
+
+def test_resolve_tools_allows_builtin_adk_tools():
+  """Verify _resolve_tools allows ADK built-in tools (no dot in name)."""
+  from google.adk.agents.llm_agent import LlmAgent
+  from google.adk.tools.tool_configs import ToolConfig
+
+  # Built-in tools have no dot — they import from google.adk.tools
+  tool_config = ToolConfig(name="google_search")
+  # Should NOT raise — this is a safe, hardcoded import path
+  resolved = LlmAgent._resolve_tools([tool_config], "/fake/path.yaml")
+  assert len(resolved) == 1
+
+
+@pytest.mark.parametrize(
+    "blocked_ref",
+    [
+        "ftplib.FTP",
+        "smtplib.SMTP",
+        "xmlrpc.client",
+        "telnetlib.Telnet",
+        "poplib.POP3",
+        "imaplib.IMAP4",
+        "asyncio.run",
+        "pathlib.Path",
+    ],
+)
+def test_newly_blocked_network_modules_are_rejected(blocked_ref: str):
+  """Verify newly added network-capable modules are blocked.
+
+  resolve_fully_qualified_name wraps errors, so we check the cause.
+  """
+  with pytest.raises(
+      ValueError, match="Invalid fully qualified name"
+  ) as exc_info:
+    config_agent_utils.resolve_fully_qualified_name(blocked_ref)
+  assert "Blocked module reference" in str(exc_info.value.__cause__)
+
+
+def test_denylist_can_be_disabled():
+  """Verify _set_enforce_denylist(False) disables module blocking."""
+  config_agent_utils._set_enforce_denylist(False)
+  try:
+    # os.getcwd is a real, importable reference — should succeed
+    result = config_agent_utils.resolve_fully_qualified_name("os.getcwd")
+    assert callable(result)
+  finally:
+    config_agent_utils._set_enforce_denylist(True)
+
+
+def test_load_config_from_path_blocks_args_when_enforced(tmp_path: Path):
+  """_load_config_from_path blocks the 'args' key when enforcement is on."""
+  config_file = tmp_path / "agent.yaml"
+  config_file.write_text("name: my_agent\nargs:\n  key: value\n")
+  config_agent_utils._set_enforce_yaml_key_denylist(True)
+  try:
+    with pytest.raises(ValueError) as exc_info:
+      config_agent_utils._load_config_from_path(str(config_file))
+    assert "Blocked key 'args' found" in str(exc_info.value)
+  finally:
+    config_agent_utils._set_enforce_yaml_key_denylist(False)

tests/unittests/agents/test_callback_context.py

@@ -0,0 +1,509 @@
+# 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 CallbackContext class."""
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import Mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.memory.memory_entry import MemoryEntry
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from google.genai.types import Part
+import pytest
+
+
+@pytest.fixture
+def mock_invocation_context():
+  """Create a mock invocation context for testing."""
+  mock_context = MagicMock()
+  mock_context.invocation_id = "test-invocation-id"
+  mock_context.agent.name = "test-agent-name"
+  mock_context.session.state = {"key1": "value1", "key2": "value2"}
+  mock_context.session.id = "test-session-id"
+  mock_context.app_name = "test-app"
+  mock_context.user_id = "test-user"
+  mock_context.artifact_service = None
+  mock_context.credential_service = None
+  return mock_context
+
+
+@pytest.fixture
+def mock_artifact_service():
+  """Create a mock artifact service for testing."""
+  mock_service = AsyncMock()
+  mock_service.list_artifact_keys.return_value = [
+      "file1.txt",
+      "file2.txt",
+      "file3.txt",
+  ]
+  return mock_service
+
+
+@pytest.fixture
+def callback_context_with_artifact_service(
+    mock_invocation_context, mock_artifact_service
+):
+  """Create a CallbackContext with a mock artifact service."""
+  mock_invocation_context.artifact_service = mock_artifact_service
+  return CallbackContext(mock_invocation_context)
+
+
+@pytest.fixture
+def callback_context_without_artifact_service(mock_invocation_context):
+  """Create a CallbackContext without an artifact service."""
+  mock_invocation_context.artifact_service = None
+  return CallbackContext(mock_invocation_context)
+
+
+@pytest.fixture
+def mock_auth_config():
+  """Create a mock auth config for testing."""
+  mock_config = Mock(spec=AuthConfig)
+  return mock_config
+
+
+@pytest.fixture
+def mock_auth_credential():
+  """Create a mock auth credential for testing."""
+  mock_credential = Mock(spec=AuthCredential)
+  mock_credential.auth_type = AuthCredentialTypes.OAUTH2
+  return mock_credential
+
+
+class TestCallbackContextListArtifacts:
+  """Test the list_artifacts method in CallbackContext."""
+
+  @pytest.mark.asyncio
+  async def test_list_artifacts_returns_artifact_keys(
+      self, callback_context_with_artifact_service, mock_artifact_service
+  ):
+    """Test that list_artifacts returns the artifact keys from the service."""
+    result = await callback_context_with_artifact_service.list_artifacts()
+
+    assert result == ["file1.txt", "file2.txt", "file3.txt"]
+    mock_artifact_service.list_artifact_keys.assert_called_once_with(
+        app_name="test-app",
+        user_id="test-user",
+        session_id="test-session-id",
+    )
+
+  @pytest.mark.asyncio
+  async def test_list_artifacts_returns_empty_list(
+      self, callback_context_with_artifact_service, mock_artifact_service
+  ):
+    """Test that list_artifacts returns an empty list when no artifacts exist."""
+    mock_artifact_service.list_artifact_keys.return_value = []
+
+    result = await callback_context_with_artifact_service.list_artifacts()
+
+    assert result == []
+    mock_artifact_service.list_artifact_keys.assert_called_once_with(
+        app_name="test-app",
+        user_id="test-user",
+        session_id="test-session-id",
+    )
+
+  @pytest.mark.asyncio
+  async def test_list_artifacts_raises_value_error_when_service_is_none(
+      self, callback_context_without_artifact_service
+  ):
+    """Test that list_artifacts raises ValueError when artifact service is None."""
+    with pytest.raises(
+        ValueError, match="Artifact service is not initialized."
+    ):
+      await callback_context_without_artifact_service.list_artifacts()
+
+  @pytest.mark.asyncio
+  async def test_list_artifacts_passes_through_service_exceptions(
+      self, callback_context_with_artifact_service, mock_artifact_service
+  ):
+    """Test that list_artifacts passes through exceptions from the artifact service."""
+    mock_artifact_service.list_artifact_keys.side_effect = Exception(
+        "Service error"
+    )
+
+    with pytest.raises(Exception, match="Service error"):
+      await callback_context_with_artifact_service.list_artifacts()
+
+
+class TestCallbackContext:
+  """Test suite for CallbackContext."""
+
+  @pytest.mark.asyncio
+  async def test_tool_context_inherits_list_artifacts(
+      self, mock_invocation_context, mock_artifact_service
+  ):
+    """Test that ToolContext inherits the list_artifacts method from CallbackContext."""
+    mock_invocation_context.artifact_service = mock_artifact_service
+    tool_context = ToolContext(mock_invocation_context)
+
+    result = await tool_context.list_artifacts()
+
+    assert result == ["file1.txt", "file2.txt", "file3.txt"]
+    mock_artifact_service.list_artifact_keys.assert_called_once_with(
+        app_name="test-app",
+        user_id="test-user",
+        session_id="test-session-id",
+    )
+
+  @pytest.mark.asyncio
+  async def test_tool_context_list_artifacts_raises_value_error_when_service_is_none(
+      self, mock_invocation_context
+  ):
+    """Test that ToolContext's list_artifacts raises ValueError when artifact service is None."""
+    mock_invocation_context.artifact_service = None
+    tool_context = ToolContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError, match="Artifact service is not initialized."
+    ):
+      await tool_context.list_artifacts()
+
+  def test_tool_context_has_list_artifacts_method(self):
+    """Test that ToolContext has the list_artifacts method available."""
+    assert hasattr(ToolContext, "list_artifacts")
+    assert callable(getattr(ToolContext, "list_artifacts"))
+
+  def test_callback_context_has_list_artifacts_method(self):
+    """Test that CallbackContext has the list_artifacts method available."""
+    assert hasattr(CallbackContext, "list_artifacts")
+    assert callable(getattr(CallbackContext, "list_artifacts"))
+
+  def test_tool_context_shares_same_list_artifacts_method_with_callback_context(
+      self,
+  ):
+    """Test that ToolContext and CallbackContext share the same list_artifacts method."""
+    assert ToolContext.list_artifacts is CallbackContext.list_artifacts
+
+  def test_initialization(self, mock_invocation_context):
+    """Test CallbackContext initialization."""
+    context = CallbackContext(mock_invocation_context)
+    assert context._invocation_context == mock_invocation_context
+    assert context._event_actions is not None
+    assert context._state is not None
+
+  @pytest.mark.asyncio
+  async def test_save_credential_with_service(
+      self, mock_invocation_context, mock_auth_config
+  ):
+    """Test save_credential when credential service is available."""
+    # Mock credential service
+    credential_service = AsyncMock()
+    mock_invocation_context.credential_service = credential_service
+
+    context = CallbackContext(mock_invocation_context)
+    await context.save_credential(mock_auth_config)
+
+    credential_service.save_credential.assert_called_once_with(
+        mock_auth_config, context
+    )
+
+  @pytest.mark.asyncio
+  async def test_save_credential_no_service(
+      self, mock_invocation_context, mock_auth_config
+  ):
+    """Test save_credential when credential service is not available."""
+    mock_invocation_context.credential_service = None
+
+    context = CallbackContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError, match="Credential service is not initialized"
+    ):
+      await context.save_credential(mock_auth_config)
+
+  @pytest.mark.asyncio
+  async def test_load_credential_with_service(
+      self, mock_invocation_context, mock_auth_config, mock_auth_credential
+  ):
+    """Test load_credential when credential service is available."""
+    # Mock credential service
+    credential_service = AsyncMock()
+    credential_service.load_credential.return_value = mock_auth_credential
+    mock_invocation_context.credential_service = credential_service
+
+    context = CallbackContext(mock_invocation_context)
+    result = await context.load_credential(mock_auth_config)
+
+    credential_service.load_credential.assert_called_once_with(
+        mock_auth_config, context
+    )
+    assert result == mock_auth_credential
+
+  @pytest.mark.asyncio
+  async def test_load_credential_no_service(
+      self, mock_invocation_context, mock_auth_config
+  ):
+    """Test load_credential when credential service is not available."""
+    mock_invocation_context.credential_service = None
+
+    context = CallbackContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError, match="Credential service is not initialized"
+    ):
+      await context.load_credential(mock_auth_config)
+
+  @pytest.mark.asyncio
+  async def test_load_credential_returns_none(
+      self, mock_invocation_context, mock_auth_config
+  ):
+    """Test load_credential returns None when credential not found."""
+    # Mock credential service
+    credential_service = AsyncMock()
+    credential_service.load_credential.return_value = None
+    mock_invocation_context.credential_service = credential_service
+
+    context = CallbackContext(mock_invocation_context)
+    result = await context.load_credential(mock_auth_config)
+
+    credential_service.load_credential.assert_called_once_with(
+        mock_auth_config, context
+    )
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_save_artifact_integration(self, mock_invocation_context):
+    """Test save_artifact to ensure credential methods follow same pattern."""
+    # Mock artifact service
+    artifact_service = AsyncMock()
+    artifact_service.save_artifact.return_value = 1
+    mock_invocation_context.artifact_service = artifact_service
+
+    context = CallbackContext(mock_invocation_context)
+    test_artifact = Part.from_text(text="test content")
+
+    version = await context.save_artifact("test_file.txt", test_artifact)
+
+    artifact_service.save_artifact.assert_called_once_with(
+        app_name="test-app",
+        user_id="test-user",
+        session_id="test-session-id",
+        filename="test_file.txt",
+        artifact=test_artifact,
+        custom_metadata=None,
+    )
+    assert version == 1
+
+  @pytest.mark.asyncio
+  async def test_load_artifact_integration(self, mock_invocation_context):
+    """Test load_artifact to ensure credential methods follow same pattern."""
+    # Mock artifact service
+    artifact_service = AsyncMock()
+    test_artifact = Part.from_text(text="test content")
+    artifact_service.load_artifact.return_value = test_artifact
+    mock_invocation_context.artifact_service = artifact_service
+
+    context = CallbackContext(mock_invocation_context)
+
+    result = await context.load_artifact("test_file.txt")
+
+    artifact_service.load_artifact.assert_called_once_with(
+        app_name="test-app",
+        user_id="test-user",
+        session_id="test-session-id",
+        filename="test_file.txt",
+        version=None,
+    )
+    assert result == test_artifact
+
+
+class TestCallbackContextAddSessionToMemory:
+  """Test the add_session_to_memory method in CallbackContext."""
+
+  @pytest.mark.asyncio
+  async def test_add_session_to_memory_success(self, mock_invocation_context):
+    """Test that add_session_to_memory calls the memory service correctly."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+
+    context = CallbackContext(mock_invocation_context)
+    await context.add_session_to_memory()
+
+    memory_service.add_session_to_memory.assert_called_once_with(
+        mock_invocation_context.session
+    )
+
+  @pytest.mark.asyncio
+  async def test_add_session_to_memory_no_service_raises(
+      self, mock_invocation_context
+  ):
+    """Test that add_session_to_memory raises ValueError when memory service is None."""
+    mock_invocation_context.memory_service = None
+
+    context = CallbackContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            r"Cannot add session to memory: memory service is not available\."
+        ),
+    ):
+      await context.add_session_to_memory()
+
+  @pytest.mark.asyncio
+  async def test_add_session_to_memory_passes_through_service_exceptions(
+      self, mock_invocation_context
+  ):
+    """Test that add_session_to_memory passes through exceptions from the memory service."""
+    memory_service = AsyncMock()
+    memory_service.add_session_to_memory.side_effect = Exception(
+        "Memory service error"
+    )
+    mock_invocation_context.memory_service = memory_service
+
+    context = CallbackContext(mock_invocation_context)
+
+    with pytest.raises(Exception, match="Memory service error"):
+      await context.add_session_to_memory()
+
+
+class TestCallbackContextAddEventsToMemory:
+  """Tests add_events_to_memory in CallbackContext."""
+
+  @pytest.mark.asyncio
+  async def test_add_events_to_memory_success(self, mock_invocation_context):
+    """Tests that add_events_to_memory calls the memory service correctly."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+    test_event = MagicMock()
+
+    context = CallbackContext(mock_invocation_context)
+    await context.add_events_to_memory(
+        events=[test_event],
+        custom_metadata={"ttl": "6000s"},
+    )
+
+    memory_service.add_events_to_memory.assert_called_once_with(
+        app_name=mock_invocation_context.session.app_name,
+        user_id=mock_invocation_context.session.user_id,
+        session_id=mock_invocation_context.session.id,
+        events=[test_event],
+        custom_metadata={"ttl": "6000s"},
+    )
+
+  @pytest.mark.asyncio
+  async def test_add_events_to_memory_no_service_raises(
+      self, mock_invocation_context
+  ):
+    """Tests that add_events_to_memory raises ValueError with no service."""
+    mock_invocation_context.memory_service = None
+
+    context = CallbackContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot add events to memory: memory service is not available\.",
+    ):
+      await context.add_events_to_memory(events=[MagicMock()])
+
+  @pytest.mark.asyncio
+  async def test_add_memory_forwards_metadata(self, mock_invocation_context):
+    """Tests that add_memory forwards memories and metadata."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+    memories = [
+        MemoryEntry(content=types.Content(parts=[types.Part(text="fact one")]))
+    ]
+    metadata = {"ttl": "6000s"}
+
+    context = CallbackContext(mock_invocation_context)
+    await context.add_memory(memories=memories, custom_metadata=metadata)
+
+    memory_service.add_memory.assert_called_once_with(
+        app_name=mock_invocation_context.session.app_name,
+        user_id=mock_invocation_context.session.user_id,
+        memories=memories,
+        custom_metadata=metadata,
+    )
+
+  @pytest.mark.asyncio
+  async def test_add_memory_accepts_memory_entries(
+      self, mock_invocation_context
+  ):
+    """Tests that add_memory forwards MemoryEntry inputs unchanged."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+    memory_entry = MemoryEntry(
+        content=types.Content(parts=[types.Part(text="fact one")])
+    )
+
+    context = CallbackContext(mock_invocation_context)
+    await context.add_memory(memories=[memory_entry])
+
+    memory_service.add_memory.assert_called_once_with(
+        app_name=mock_invocation_context.session.app_name,
+        user_id=mock_invocation_context.session.user_id,
+        memories=[memory_entry],
+        custom_metadata=None,
+    )
+
+  @pytest.mark.asyncio
+  async def test_add_memory_no_service_raises(self, mock_invocation_context):
+    """Tests that add_memory raises ValueError with no service."""
+    mock_invocation_context.memory_service = None
+
+    context = CallbackContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot add memory: memory service is not available\.",
+    ):
+      await context.add_memory(
+          memories=[
+              MemoryEntry(
+                  content=types.Content(parts=[types.Part(text="fact one")])
+              )
+          ]
+      )
+
+
+class TestToolContextAddSessionToMemory:
+  """Test the add_session_to_memory method in ToolContext."""
+
+  @pytest.mark.asyncio
+  async def test_add_session_to_memory_success(self, mock_invocation_context):
+    """Test that ToolContext.add_session_to_memory calls the memory service correctly."""
+    memory_service = AsyncMock()
+    mock_invocation_context.memory_service = memory_service
+
+    tool_context = ToolContext(mock_invocation_context)
+    await tool_context.add_session_to_memory()
+
+    memory_service.add_session_to_memory.assert_called_once_with(
+        mock_invocation_context.session
+    )
+
+  @pytest.mark.asyncio
+  async def test_add_session_to_memory_no_service_raises(
+      self, mock_invocation_context
+  ):
+    """Test that ToolContext.add_session_to_memory raises ValueError when memory service is None."""
+    mock_invocation_context.memory_service = None
+
+    tool_context = ToolContext(mock_invocation_context)
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            r"Cannot add session to memory: memory service is not available\."
+        ),
+    ):
+      await tool_context.add_session_to_memory()

tests/unittests/agents/test_context_cache_config.py

@@ -0,0 +1,174 @@
+# 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 ContextCacheConfig."""
+
+from google.adk.agents.context_cache_config import ContextCacheConfig
+from pydantic import ValidationError
+import pytest
+
+
+class TestContextCacheConfig:
+  """Test suite for ContextCacheConfig."""
+
+  def test_default_values(self):
+    """Test that default values are set correctly."""
+    config = ContextCacheConfig()
+
+    assert config.cache_intervals == 10
+    assert config.ttl_seconds == 1800  # 30 minutes
+    assert config.min_tokens == 0
+
+  def test_custom_values(self):
+    """Test creating config with custom values."""
+    config = ContextCacheConfig(
+        cache_intervals=15, ttl_seconds=3600, min_tokens=1024
+    )
+
+    assert config.cache_intervals == 15
+    assert config.ttl_seconds == 3600
+    assert config.min_tokens == 1024
+
+  def test_cache_intervals_validation(self):
+    """Test cache_intervals validation constraints."""
+    # Valid range
+    config = ContextCacheConfig(cache_intervals=1)
+    assert config.cache_intervals == 1
+
+    config = ContextCacheConfig(cache_intervals=100)
+    assert config.cache_intervals == 100
+
+    # Invalid: too low
+    with pytest.raises(ValidationError) as exc_info:
+      ContextCacheConfig(cache_intervals=0)
+    assert "greater than or equal to 1" in str(exc_info.value)
+
+    # Invalid: too high
+    with pytest.raises(ValidationError) as exc_info:
+      ContextCacheConfig(cache_intervals=101)
+    assert "less than or equal to 100" in str(exc_info.value)
+
+  def test_ttl_seconds_validation(self):
+    """Test ttl_seconds validation constraints."""
+    # Valid range
+    config = ContextCacheConfig(ttl_seconds=1)
+    assert config.ttl_seconds == 1
+
+    config = ContextCacheConfig(ttl_seconds=86400)  # 24 hours
+    assert config.ttl_seconds == 86400
+
+    # Invalid: zero or negative
+    with pytest.raises(ValidationError) as exc_info:
+      ContextCacheConfig(ttl_seconds=0)
+    assert "greater than 0" in str(exc_info.value)
+
+    with pytest.raises(ValidationError) as exc_info:
+      ContextCacheConfig(ttl_seconds=-1)
+    assert "greater than 0" in str(exc_info.value)
+
+  def test_min_tokens_validation(self):
+    """Test min_tokens validation constraints."""
+    # Valid values
+    config = ContextCacheConfig(min_tokens=0)
+    assert config.min_tokens == 0
+
+    config = ContextCacheConfig(min_tokens=1024)
+    assert config.min_tokens == 1024
+
+    # Invalid: negative
+    with pytest.raises(ValidationError) as exc_info:
+      ContextCacheConfig(min_tokens=-1)
+    assert "greater than or equal to 0" in str(exc_info.value)
+
+  def test_ttl_string_property(self):
+    """Test ttl_string property returns correct format."""
+    config = ContextCacheConfig(ttl_seconds=1800)
+    assert config.ttl_string == "1800s"
+
+    config = ContextCacheConfig(ttl_seconds=3600)
+    assert config.ttl_string == "3600s"
+
+  def test_str_representation(self):
+    """Test string representation for logging."""
+    config = ContextCacheConfig(
+        cache_intervals=15, ttl_seconds=3600, min_tokens=1024
+    )
+
+    expected = (
+        "ContextCacheConfig(cache_intervals=15, ttl=3600s, min_tokens=1024, "
+        "create_http_options=None)"
+    )
+    assert str(config) == expected
+
+  def test_str_representation_defaults(self):
+    """Test string representation with default values."""
+    config = ContextCacheConfig()
+
+    expected = (
+        "ContextCacheConfig(cache_intervals=10, ttl=1800s, min_tokens=0, "
+        "create_http_options=None)"
+    )
+    assert str(config) == expected
+
+  def test_pydantic_model_validation(self):
+    """Test that Pydantic model validation works correctly."""
+    # Test extra fields are forbidden
+    with pytest.raises(ValidationError) as exc_info:
+      ContextCacheConfig(cache_intervals=10, extra_field="not_allowed")
+    assert "extra" in str(exc_info.value).lower()
+
+  def test_field_descriptions(self):
+    """Test that fields have proper descriptions."""
+    fields = ContextCacheConfig.model_fields
+
+    assert "cache_intervals" in fields
+    assert (
+        "Maximum number of invocations" in fields["cache_intervals"].description
+    )
+
+    assert "ttl_seconds" in fields
+    assert "Time-to-live for cache" in fields["ttl_seconds"].description
+
+    assert "min_tokens" in fields
+    assert "Minimum prior-request tokens" in fields["min_tokens"].description
+
+  def test_immutability_config(self):
+    """Test that the model config is set correctly."""
+    config = ContextCacheConfig()
+    assert config.model_config["extra"] == "forbid"
+
+  def test_realistic_scenarios(self):
+    """Test realistic configuration scenarios."""
+    # Quick caching for development
+    dev_config = ContextCacheConfig(
+        cache_intervals=5, ttl_seconds=600, min_tokens=0  # 10 minutes
+    )
+    assert dev_config.cache_intervals == 5
+    assert dev_config.ttl_seconds == 600
+
+    # Production caching
+    prod_config = ContextCacheConfig(
+        cache_intervals=20, ttl_seconds=7200, min_tokens=2048  # 2 hours
+    )
+    assert prod_config.cache_intervals == 20
+    assert prod_config.ttl_seconds == 7200
+    assert prod_config.min_tokens == 2048
+
+    # Conservative caching
+    conservative_config = ContextCacheConfig(
+        cache_intervals=3, ttl_seconds=300, min_tokens=4096  # 5 minutes
+    )
+    assert conservative_config.cache_intervals == 3
+    assert conservative_config.ttl_seconds == 300
+    assert conservative_config.min_tokens == 4096

tests/unittests/agents/test_invocation_context.py

@@ -0,0 +1,694 @@
+# 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.
+
+from unittest.mock import Mock
+
+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.apps import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.events.event_actions import EventActions
+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
+
+from .. import testing_utils
+
+
+class TestInvocationContext:
+  """Test suite for InvocationContext."""
+
+  @pytest.fixture
+  def mock_events(self):
+    """Create mock events for testing."""
+    event1 = Mock(spec=Event)
+    event1.invocation_id = 'inv_1'
+    event1.branch = 'agent_1'
+
+    event2 = Mock(spec=Event)
+    event2.invocation_id = 'inv_1'
+    event2.branch = 'agent_2'
+
+    event3 = Mock(spec=Event)
+    event3.invocation_id = 'inv_2'
+    event3.branch = 'agent_1'
+
+    event4 = Mock(spec=Event)
+    event4.invocation_id = 'inv_2'
+    event4.branch = 'agent_2'
+
+    return [event1, event2, event3, event4]
+
+  @pytest.fixture
+  def mock_invocation_context(self, mock_events):
+    """Create a mock invocation context for testing."""
+    ctx = InvocationContext(
+        session_service=Mock(spec=BaseSessionService),
+        agent=Mock(spec=BaseAgent),
+        invocation_id='inv_1',
+        branch='agent_1',
+        session=Mock(spec=Session, events=mock_events),
+    )
+    return ctx
+
+  def test_get_events_returns_all_events_by_default(
+      self, mock_invocation_context, mock_events
+  ):
+    """Tests that get_events returns all events when no filters are applied."""
+    events = mock_invocation_context._get_events()
+    assert events == mock_events
+
+  def test_get_events_filters_by_current_invocation(
+      self, mock_invocation_context, mock_events
+  ):
+    """Tests that get_events correctly filters by the current invocation."""
+    event1, event2, _, _ = mock_events
+    events = mock_invocation_context._get_events(current_invocation=True)
+    assert events == [event1, event2]
+
+  def test_get_events_filters_by_current_branch(
+      self, mock_invocation_context, mock_events
+  ):
+    """Tests that get_events correctly filters by the current branch."""
+    event1, _, event3, _ = mock_events
+    events = mock_invocation_context._get_events(current_branch=True)
+    assert events == [event1, event3]
+
+  def test_get_events_filters_by_invocation_and_branch(
+      self, mock_invocation_context, mock_events
+  ):
+    """Tests that get_events filters by invocation and branch."""
+    event1, _, _, _ = mock_events
+    events = mock_invocation_context._get_events(
+        current_invocation=True,
+        current_branch=True,
+    )
+    assert events == [event1]
+
+  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()
+    assert not events
+
+  def test_get_events_with_no_matching_events(self, mock_invocation_context):
+    """Tests get_events when no events match the filters."""
+    mock_invocation_context.invocation_id = 'inv_3'
+    mock_invocation_context.branch = 'branch_C'
+
+    # Filter by invocation
+    events = mock_invocation_context._get_events(current_invocation=True)
+    assert not events
+
+    # Filter by branch
+    events = mock_invocation_context._get_events(current_branch=True)
+    assert not events
+
+    # Filter by both
+    events = mock_invocation_context._get_events(
+        current_invocation=True,
+        current_branch=True,
+    )
+    assert not events
+
+
+class TestInvocationContextWithAppResumablity:
+  """Test suite for InvocationContext regarding app resumability."""
+
+  @pytest.fixture
+  def long_running_function_call(self) -> FunctionCall:
+    """A long running function call."""
+    return FunctionCall(
+        id='tool_call_id_1',
+        name='long_running_function_call',
+        args={},
+    )
+
+  @pytest.fixture
+  def event_to_pause(self, long_running_function_call) -> Event:
+    """An event with a long running function call."""
+    return Event(
+        invocation_id='inv_1',
+        author='agent',
+        content=testing_utils.ModelContent(
+            [Part(function_call=long_running_function_call)]
+        ),
+        long_running_tool_ids=[long_running_function_call.id],
+    )
+
+  def _create_test_invocation_context(
+      self, resumability_config: ResumabilityConfig | None = None
+  ) -> InvocationContext:
+    """Create a mock invocation context for testing."""
+    ctx = InvocationContext(
+        session_service=Mock(spec=BaseSessionService),
+        agent=Mock(spec=BaseAgent),
+        invocation_id='inv_1',
+        session=Mock(spec=Session, events=[]),
+        resumability_config=resumability_config,
+    )
+    return ctx
+
+  def test_should_pause_invocation_with_resumable_app(self, event_to_pause):
+    """Tests should_pause_invocation with a resumable app."""
+    mock_invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+
+    assert mock_invocation_context.should_pause_invocation(event_to_pause)
+
+  def test_should_pause_invocation_with_non_resumable_app(self, event_to_pause):
+    """Tests should_pause_invocation pauses even without resumability."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=False)
+    )
+
+    assert invocation_context.should_pause_invocation(event_to_pause)
+
+  def test_should_not_pause_invocation_with_no_long_running_tool_ids(
+      self, event_to_pause
+  ):
+    """Tests should_pause_invocation with no long running tools."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    nonpausable_event = event_to_pause.model_copy(
+        update={'long_running_tool_ids': []}
+    )
+
+    assert not invocation_context.should_pause_invocation(nonpausable_event)
+
+  def test_should_not_pause_invocation_with_no_function_calls(
+      self, event_to_pause
+  ):
+    """Tests should_pause_invocation with a non-model event."""
+    mock_invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    nonpausable_event = event_to_pause.model_copy(
+        update={'content': testing_utils.UserContent('test text part')}
+    )
+
+    assert not mock_invocation_context.should_pause_invocation(
+        nonpausable_event
+    )
+
+  def test_should_not_pause_when_user_resumes_in_sub_branch(
+      self, event_to_pause, long_running_function_call
+  ):
+    """We do not pause the invocation if a subsequent user event belongs to a sub-branch."""
+    # Arrange
+    mock_invocation_context = self._create_test_invocation_context()
+    user_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        branch=f'agent@{long_running_function_call.id}.child',
+    )
+    mock_invocation_context.session.events = [event_to_pause, user_event]
+
+    # Act
+    should_pause = mock_invocation_context.should_pause_invocation(
+        event_to_pause
+    )
+
+    # Assert
+    assert not should_pause
+
+  def test_should_not_pause_when_user_resumes_in_deeply_nested_sub_branch(
+      self, event_to_pause, long_running_function_call
+  ):
+    """We do not pause if the user resumes in a deeply nested sub-branch containing the tool call."""
+    # Arrange
+    mock_invocation_context = self._create_test_invocation_context()
+    user_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        branch=f'parent@other.child@{long_running_function_call.id}.grandchild',
+    )
+    mock_invocation_context.session.events = [event_to_pause, user_event]
+
+    # Act
+    should_pause = mock_invocation_context.should_pause_invocation(
+        event_to_pause
+    )
+
+    # Assert
+    assert not should_pause
+
+  def test_should_pause_when_user_resumes_in_different_branch(
+      self, event_to_pause
+  ):
+    """We still pause the invocation if the subsequent user event belongs to a different branch."""
+    # Arrange
+    mock_invocation_context = self._create_test_invocation_context()
+    user_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        branch='parent@different_id.child',
+    )
+    mock_invocation_context.session.events = [event_to_pause, user_event]
+
+    # Act
+    should_pause = mock_invocation_context.should_pause_invocation(
+        event_to_pause
+    )
+
+    # Assert
+    assert should_pause
+
+  def test_is_resumable_true(self):
+    """Tests that is_resumable is True when resumability is enabled."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    assert invocation_context.is_resumable
+
+  def test_is_resumable_false(self):
+    """Tests that is_resumable is False when resumability is disabled."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=False)
+    )
+    assert not invocation_context.is_resumable
+
+  def test_is_resumable_no_config(self):
+    """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
+
+  def test_populate_invocation_agent_states_not_resumable(self):
+    """Tests that populate_invocation_agent_states does nothing if not resumable."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=False)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=True, agent_state=None),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert not invocation_context.end_of_agents
+
+  def test_populate_invocation_agent_states_end_of_agent(self):
+    """Tests that populate_invocation_agent_states handles end_of_agent."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=True, agent_state=None),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert invocation_context.end_of_agents == {'agent1': True}
+
+  def test_populate_invocation_agent_states_with_agent_state(self):
+    """Tests that populate_invocation_agent_states handles agent_state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(
+            end_of_agent=False,
+            agent_state=BaseAgentState().model_dump(mode='json'),
+        ),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert invocation_context.agent_states == {'agent1': {}}
+    assert invocation_context.end_of_agents == {'agent1': False}
+
+  def test_populate_invocation_agent_states_with_agent_state_and_end_of_agent(
+      self,
+  ):
+    """Tests that populate_invocation_agent_states handles agent_state and end_of_agent."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(
+            end_of_agent=True,
+            agent_state=BaseAgentState().model_dump(mode='json'),
+        ),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    # When both agent_state and end_of_agent are set, agent_state should be
+    # cleared, as end_of_agent is of a higher priority.
+    assert not invocation_context.agent_states
+    assert invocation_context.end_of_agents == {'agent1': True}
+
+  def test_populate_invocation_agent_states_with_content_no_state(self):
+    """Tests that populate_invocation_agent_states creates default state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=False, agent_state=None),
+        content=Content(role='model', parts=[Part(text='hi')]),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert invocation_context.agent_states == {
+        'agent1': BaseAgentState().model_dump(mode='json')
+    }
+    assert invocation_context.end_of_agents == {'agent1': False}
+
+  def test_populate_invocation_agent_states_user_message_event(self):
+    """Tests that populate_invocation_agent_states ignores user message events for default state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='user',
+        actions=EventActions(end_of_agent=False, agent_state=None),
+        content=Content(role='user', parts=[Part(text='hi')]),
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert not invocation_context.end_of_agents
+
+  def test_populate_invocation_agent_states_no_content(self):
+    """Tests that populate_invocation_agent_states ignores events with no content if no state."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    event = Event(
+        invocation_id='inv_1',
+        author='agent1',
+        actions=EventActions(end_of_agent=None, agent_state=None),
+        content=None,
+    )
+    invocation_context.session.events = [event]
+    invocation_context.populate_invocation_agent_states()
+    assert not invocation_context.agent_states
+    assert not invocation_context.end_of_agents
+
+  def test_set_agent_state_with_end_of_agent_true(self):
+    """Tests that set_agent_state clears agent_state and sets end_of_agent to True."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    invocation_context.agent_states['agent1'] = {}
+    invocation_context.end_of_agents['agent1'] = False
+
+    # Set state with end_of_agent=True, which should clear the existing
+    # agent_state.
+    invocation_context.set_agent_state('agent1', end_of_agent=True)
+    assert 'agent1' not in invocation_context.agent_states
+    assert invocation_context.end_of_agents['agent1']
+
+  def test_set_agent_state_with_agent_state(self):
+    """Tests that set_agent_state sets agent_state and sets end_of_agent to False."""
+    agent_state = BaseAgentState()
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    invocation_context.end_of_agents['agent1'] = True
+
+    # Set state with agent_state=agent_state, which should set the agent_state
+    # and reset the end_of_agent flag to False.
+    invocation_context.set_agent_state('agent1', agent_state=agent_state)
+    assert invocation_context.agent_states['agent1'] == agent_state.model_dump(
+        mode='json'
+    )
+    assert invocation_context.end_of_agents['agent1'] is False
+
+  def test_reset_agent_state(self):
+    """Tests that set_agent_state clears agent_state and end_of_agent."""
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    invocation_context.agent_states['agent1'] = {}
+    invocation_context.end_of_agents['agent1'] = True
+
+    # Reset state, which should clear the agent_state and end_of_agent flag.
+    invocation_context.set_agent_state('agent1')
+    assert 'agent1' not in invocation_context.agent_states
+    assert 'agent1' not in invocation_context.end_of_agents
+
+  def test_reset_sub_agent_states(self):
+    """Tests that reset_sub_agent_states resets sub-agent states."""
+    sub_sub_agent_1 = BaseAgent(name='sub_sub_agent_1')
+    sub_agent_1 = BaseAgent(name='sub_agent_1', sub_agents=[sub_sub_agent_1])
+    sub_agent_2 = BaseAgent(name='sub_agent_2')
+    root_agent = BaseAgent(
+        name='root_agent', sub_agents=[sub_agent_1, sub_agent_2]
+    )
+
+    invocation_context = self._create_test_invocation_context(
+        ResumabilityConfig(is_resumable=True)
+    )
+    invocation_context.agent = root_agent
+    invocation_context.set_agent_state(
+        'sub_agent_1', agent_state=BaseAgentState()
+    )
+    invocation_context.set_agent_state('sub_agent_2', end_of_agent=True)
+    invocation_context.set_agent_state(
+        'sub_sub_agent_1', agent_state=BaseAgentState()
+    )
+
+    assert 'sub_agent_1' in invocation_context.agent_states
+    assert 'sub_agent_2' in invocation_context.end_of_agents
+    assert 'sub_sub_agent_1' in invocation_context.agent_states
+
+    invocation_context.reset_sub_agent_states('root_agent')
+
+    assert 'sub_agent_1' not in invocation_context.agent_states
+    assert 'sub_agent_1' not in invocation_context.end_of_agents
+    assert 'sub_agent_2' not in invocation_context.agent_states
+    assert 'sub_agent_2' not in invocation_context.end_of_agents
+    assert 'sub_sub_agent_1' not in invocation_context.agent_states
+    assert 'sub_sub_agent_1' not in invocation_context.end_of_agents
+
+
+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
+
+  def test_stamp_event_branch_context_preserves_isolation_scope(
+      self, test_invocation_context
+  ):
+    """Tests stamp_event_branch_context does not overwrite existing isolation_scope with None."""
+    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',
+        branch='root@1',
+        isolation_scope=None,  # Coordinator FC has None scope
+        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',
+        isolation_scope='task_123',  # Pre-populated active task scope
+        content=Content(role='user', parts=[fr]),
+    )
+    invocation_context = test_invocation_context([fc_event, fr_event])
+
+    invocation_context.stamp_event_branch_context(fr_event)
+    assert fr_event.branch == 'root@1'
+    assert fr_event.isolation_scope == 'task_123'
+
+  def test_stamp_event_branch_context_does_not_overwrite_existing_scope(
+      self, test_invocation_context
+  ):
+    """Tests stamp_event_branch_context does not overwrite existing isolation_scope if set."""
+    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',
+        branch='root@1',
+        isolation_scope='task_456',  # Function call has isolation scope
+        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',
+        isolation_scope='task_123',  # Pre-populated active task scope
+        content=Content(role='user', parts=[fr]),
+    )
+    invocation_context = test_invocation_context([fc_event, fr_event])
+
+    invocation_context.stamp_event_branch_context(fr_event)
+    assert fr_event.branch == 'root@1'
+    assert fr_event.isolation_scope == 'task_123'
+
+  def test_find_matching_function_call_when_response_is_not_last_event(
+      self, test_invocation_context
+  ):
+    """Tests that matching function call is found even when response is not the last event in history."""
+    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]),
+    )
+    # Add a subsequent event to the history so that fr_event is NOT the last one
+    subsequent_event = Event(
+        invocation_id='inv_1',
+        author='user',
+        content=Content(role='user', parts=[Part(text='next user message')]),
+    )
+    invocation_context = test_invocation_context(
+        [fc_event, fr_event, subsequent_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)

tests/unittests/agents/test_invocation_context_process_queue.py

@@ -0,0 +1,159 @@
+# 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 _event_queue and _enqueue_event on InvocationContext."""
+
+from __future__ import annotations
+
+import asyncio
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.events.event import Event
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+import pytest
+
+
+async def _create_ic_with_queue() -> InvocationContext:
+  """Create a minimal InvocationContext with _event_queue set."""
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      instruction='test',
+  )
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  ic = InvocationContext(
+      invocation_id='test_invocation',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  ic._event_queue = asyncio.Queue()
+  return ic
+
+
+async def _create_ic_without_queue() -> InvocationContext:
+  """Create a minimal InvocationContext without _event_queue."""
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      instruction='test',
+  )
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id='test_invocation',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+
+
+@pytest.mark.asyncio
+async def test_non_partial_event_blocks_until_processed() -> None:
+  """A non-partial event should block _enqueue_event until the consumer
+  signals processed."""
+  ic: InvocationContext = await _create_ic_with_queue()
+  event: Event = Event(id=Event.new_id(), author='test')
+
+  completed: bool = False
+
+  async def consumer() -> None:
+    nonlocal completed
+    ev, processed = await ic._event_queue.get()
+    assert ev is event
+    assert processed is not None
+    processed.set()
+    completed = True
+
+  consumer_task: asyncio.Task = asyncio.create_task(consumer())
+  await ic._enqueue_event(event)
+  await consumer_task
+
+  assert completed, 'Consumer should have processed the event.'
+
+
+@pytest.mark.asyncio
+async def test_partial_event_does_not_block() -> None:
+  """A partial event should not block — it returns immediately
+  without waiting for a processed signal."""
+  ic: InvocationContext = await _create_ic_with_queue()
+  event: Event = Event(id=Event.new_id(), author='test', partial=True)
+
+  await ic._enqueue_event(event)
+
+  assert not ic._event_queue.empty()
+  ev, processed = await ic._event_queue.get()
+  assert ev is event
+  assert processed is None, 'Partial events should have no processed signal.'
+
+
+@pytest.mark.asyncio
+async def test_events_arrive_in_order() -> None:
+  """Multiple partial events should arrive on the queue in order."""
+  ic: InvocationContext = await _create_ic_with_queue()
+  events: list[Event] = [
+      Event(id=Event.new_id(), author=f'test_{i}', partial=True)
+      for i in range(5)
+  ]
+
+  for event in events:
+    await ic._enqueue_event(event)
+
+  for i in range(5):
+    ev, _ = await ic._event_queue.get()
+    assert ev.author == f'test_{i}'
+
+
+@pytest.mark.asyncio
+async def test_enqueue_event_raises_when_queue_not_set() -> None:
+  """_enqueue_event should raise RuntimeError if _event_queue is None."""
+  ic: InvocationContext = await _create_ic_without_queue()
+  event: Event = Event(id=Event.new_id(), author='test')
+
+  with pytest.raises(RuntimeError, match='_event_queue is not set'):
+    await ic._enqueue_event(event)
+
+
+@pytest.mark.asyncio
+async def test_non_partial_event_waits_for_signal() -> None:
+  """Verify that _enqueue_event for a non-partial event actually waits —
+  it should not complete before the consumer signals."""
+  ic: InvocationContext = await _create_ic_with_queue()
+  event: Event = Event(id=Event.new_id(), author='test')
+
+  emit_done: bool = False
+
+  async def emitter() -> None:
+    nonlocal emit_done
+    await ic._enqueue_event(event)
+    emit_done = True
+
+  emit_task: asyncio.Task = asyncio.create_task(emitter())
+
+  # Give the emitter a chance to run.
+  await asyncio.sleep(0.01)
+  assert not emit_done, '_enqueue_event should still be waiting.'
+
+  # Now consume and signal.
+  _, processed = await ic._event_queue.get()
+  processed.set()
+
+  await emit_task
+  assert emit_done, '_enqueue_event should complete after signal.'

tests/unittests/agents/test_langgraph_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.
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Skip all tests in this module if LangGraph dependencies are not available
+LANGGRAPH_AVAILABLE = True
+try:
+  from google.adk.agents.invocation_context import InvocationContext
+  from google.adk.agents.langgraph_agent import LangGraphAgent
+  from google.adk.events.event import Event
+  from google.adk.plugins.plugin_manager import PluginManager
+  from google.genai import types
+  from langchain_core.messages import AIMessage
+  from langchain_core.messages import HumanMessage
+  from langchain_core.messages import SystemMessage
+  from langgraph.graph.graph import CompiledGraph
+except ImportError:
+  LANGGRAPH_AVAILABLE = False
+
+  # IMPORTANT: Dummy classes are REQUIRED in this file but NOT in A2A test files.
+  # Here's why this file is different from A2A test files:
+  #
+  # 1. MODULE-LEVEL USAGE IN DECORATORS:
+  #    This file uses @pytest.mark.parametrize decorator with complex nested structures
+  #    that directly reference imported types like Event(), types.Content(), types.Part.from_text().
+  #    These decorator expressions are evaluated during MODULE COMPILATION TIME,
+  #    not during test execution time.
+  #
+  # 2. A2A TEST FILES PATTERN:
+  #    Most A2A test files only use imported types within test method bodies:
+  #    - Inside test functions: def test_something(): Message(...)
+  #    - These are evaluated during TEST EXECUTION TIME when tests are skipped
+  #    - No NameError occurs because skipped tests don't execute their bodies
+  #
+  # 3. WHAT HAPPENS WITHOUT DUMMIES:
+  #    If we remove dummy classes from this file:
+  #    - Python tries to compile the @pytest.mark.parametrize decorator
+  #    - It encounters Event(...), types.Content(...), etc.
+  #    - These names are undefined → NameError during module compilation
+  #    - Test collection fails before pytest.mark.skipif can even run
+  #
+  # 4. WHY DUMMIES WORK:
+  #    - DummyTypes() can be called like Event() → returns DummyTypes instance
+  #    - DummyTypes.__getattr__ handles types.Content → returns DummyTypes instance
+  #    - DummyTypes.__call__ handles types.Part.from_text() → returns DummyTypes instance
+  #    - The parametrize decorator gets dummy objects instead of real ones
+  #    - Tests still get skipped due to pytestmark, so dummies never actually run
+  #
+  # 5. EXCEPTION CASES IN A2A FILES:
+  #    A few A2A files DID need dummies initially because they had:
+  #    - Type annotations: def create_helper(x: str) -> Message
+  #    - But we removed those type annotations to eliminate the need for dummies
+  #
+  # This file cannot avoid dummies because the parametrize decorator usage
+  # is fundamental to the test structure and cannot be easily refactored.
+
+  class DummyTypes:
+
+    def __getattr__(self, name):
+      return DummyTypes()
+
+    def __call__(self, *args, **kwargs):
+      return DummyTypes()
+
+  InvocationContext = DummyTypes()
+  LangGraphAgent = DummyTypes()
+  Event = DummyTypes()
+  PluginManager = DummyTypes()
+  types = (
+      DummyTypes()
+  )  # Must support chained calls like types.Content(), types.Part.from_text()
+  AIMessage = DummyTypes()
+  HumanMessage = DummyTypes()
+  SystemMessage = DummyTypes()
+  CompiledGraph = DummyTypes()
+
+pytestmark = pytest.mark.skipif(
+    not LANGGRAPH_AVAILABLE, reason="LangGraph dependencies not available"
+)
+
+
+@pytest.mark.parametrize(
+    "checkpointer_value, events_list, expected_messages",
+    [
+        (
+            MagicMock(),
+            [
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="user",
+                    content=types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text="test prompt")],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="root_agent",
+                    content=types.Content(
+                        role="model",
+                        parts=[types.Part.from_text(text="(some delegation)")],
+                    ),
+                ),
+            ],
+            [
+                SystemMessage(content="test system prompt"),
+                HumanMessage(content="test prompt"),
+            ],
+        ),
+        (
+            None,
+            [
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="user",
+                    content=types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text="user prompt 1")],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="root_agent",
+                    content=types.Content(
+                        role="model",
+                        parts=[
+                            types.Part.from_text(text="root agent response")
+                        ],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="weather_agent",
+                    content=types.Content(
+                        role="model",
+                        parts=[
+                            types.Part.from_text(text="weather agent response")
+                        ],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="user",
+                    content=types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text="user prompt 2")],
+                    ),
+                ),
+            ],
+            [
+                SystemMessage(content="test system prompt"),
+                HumanMessage(content="user prompt 1"),
+                AIMessage(content="weather agent response"),
+                HumanMessage(content="user prompt 2"),
+            ],
+        ),
+        (
+            MagicMock(),
+            [
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="user",
+                    content=types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text="user prompt 1")],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="root_agent",
+                    content=types.Content(
+                        role="model",
+                        parts=[
+                            types.Part.from_text(text="root agent response")
+                        ],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="weather_agent",
+                    content=types.Content(
+                        role="model",
+                        parts=[
+                            types.Part.from_text(text="weather agent response")
+                        ],
+                    ),
+                ),
+                Event(
+                    invocation_id="test_invocation_id",
+                    author="user",
+                    content=types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text="user prompt 2")],
+                    ),
+                ),
+            ],
+            [
+                SystemMessage(content="test system prompt"),
+                HumanMessage(content="user prompt 2"),
+            ],
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_langgraph_agent(
+    checkpointer_value, events_list, expected_messages
+):
+  mock_graph = MagicMock(spec=CompiledGraph)
+  mock_graph_state = MagicMock()
+  mock_graph_state.values = {}
+  mock_graph.get_state.return_value = mock_graph_state
+
+  mock_graph.checkpointer = checkpointer_value
+  mock_graph.invoke.return_value = {
+      "messages": [AIMessage(content="test response")]
+  }
+
+  mock_parent_context = MagicMock(spec=InvocationContext)
+  mock_parent_context._state_schema = None
+  mock_session = MagicMock()
+  mock_parent_context.session = mock_session
+  mock_parent_context.user_content = types.Content(
+      role="user", parts=[types.Part.from_text(text="test prompt")]
+  )
+  mock_parent_context.branch = "parent_agent"
+  mock_parent_context.end_invocation = False
+  mock_session.events = events_list
+  mock_parent_context.invocation_id = "test_invocation_id"
+  mock_parent_context.model_copy.return_value = mock_parent_context
+  mock_parent_context.plugin_manager = PluginManager(plugins=[])
+
+  weather_agent = LangGraphAgent(
+      name="weather_agent",
+      description="A agent that answers weather questions",
+      instruction="test system prompt",
+      graph=mock_graph,
+  )
+
+  result_event = None
+  async for event in weather_agent.run_async(mock_parent_context):
+    result_event = event
+
+  assert result_event.author == "weather_agent"
+  assert result_event.content.parts[0].text == "test response"
+
+  mock_graph.invoke.assert_called_once()
+  mock_graph.invoke.assert_called_with(
+      {"messages": expected_messages},
+      {"configurable": {"thread_id": mock_session.id}},
+  )

tests/unittests/agents/test_live_request_queue.py

@@ -0,0 +1,81 @@
+# 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.
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.agents.live_request_queue import LiveRequest
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.genai import types
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_close_queue():
+  queue = LiveRequestQueue()
+
+  with patch.object(queue._queue, "put_nowait") as mock_put_nowait:
+    queue.close()
+    mock_put_nowait.assert_called_once_with(LiveRequest(close=True))
+
+
+def test_send_content():
+  queue = LiveRequestQueue()
+  content = MagicMock(spec=types.Content)
+
+  with patch.object(queue._queue, "put_nowait") as mock_put_nowait:
+    queue.send_content(content)
+    mock_put_nowait.assert_called_once_with(LiveRequest(content=content))
+
+
+def test_send_content_sets_partial():
+  queue = LiveRequestQueue()
+  content = MagicMock(spec=types.Content)
+
+  with patch.object(queue._queue, "put_nowait") as mock_put_nowait:
+    queue.send_content(content, partial=True)
+    mock_put_nowait.assert_called_once_with(
+        LiveRequest(content=content, partial=True)
+    )
+
+
+def test_send_realtime():
+  queue = LiveRequestQueue()
+  blob = MagicMock(spec=types.Blob)
+
+  with patch.object(queue._queue, "put_nowait") as mock_put_nowait:
+    queue.send_realtime(blob)
+    mock_put_nowait.assert_called_once_with(LiveRequest(blob=blob))
+
+
+def test_send():
+  queue = LiveRequestQueue()
+  req = LiveRequest(content=MagicMock(spec=types.Content))
+
+  with patch.object(queue._queue, "put_nowait") as mock_put_nowait:
+    queue.send(req)
+    mock_put_nowait.assert_called_once_with(req)
+
+
+@pytest.mark.asyncio
+async def test_get():
+  queue = LiveRequestQueue()
+  res = MagicMock(spec=types.Content)
+
+  with patch.object(queue._queue, "get", return_value=res) as mock_get:
+    result = await queue.get()
+
+    assert result == res
+    mock_get.assert_called_once()

tests/unittests/agents/test_llm_agent_callbacks.py

@@ -0,0 +1,210 @@
+# 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.
+
+from typing import Any
+from typing import Optional
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from .. import testing_utils
+
+
+class MockBeforeModelCallback(BaseModel):
+  mock_response: str
+
+  def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+class MockAfterModelCallback(BaseModel):
+  mock_response: str
+
+  def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_response: LlmResponse,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+class MockAsyncBeforeModelCallback(BaseModel):
+  mock_response: str
+
+  async def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+class MockAsyncAfterModelCallback(BaseModel):
+  mock_response: str
+
+  async def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_response: LlmResponse,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+def noop_callback(**kwargs) -> Optional[LlmResponse]:
+  pass
+
+
+async def async_noop_callback(**kwargs) -> Optional[LlmResponse]:
+  pass
+
+
+@pytest.mark.asyncio
+async def test_before_model_callback():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=MockBeforeModelCallback(
+          mock_response='before_model_callback'
+      ),
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'before_model_callback'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_before_model_callback_noop():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=noop_callback,
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_after_model_callback():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_model_callback=MockAfterModelCallback(
+          mock_response='after_model_callback'
+      ),
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'after_model_callback'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_async_before_model_callback():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=MockAsyncBeforeModelCallback(
+          mock_response='async_before_model_callback'
+      ),
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'async_before_model_callback'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_async_before_model_callback_noop():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=async_noop_callback,
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_async_after_model_callback():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_model_callback=MockAsyncAfterModelCallback(
+          mock_response='async_after_model_callback'
+      ),
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'async_after_model_callback'),
+  ]

tests/unittests/agents/test_llm_agent_error_messages.py

@@ -0,0 +1,90 @@
+# 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 enhanced error messages in agent handling."""
+
+from google.adk.agents.llm_agent import LlmAgent
+import pytest
+
+
+def test_agent_not_found_enhanced_error():
+  """Verify enhanced error message for agent not found."""
+  root_agent = LlmAgent(
+      name='root',
+      model='gemini-2.5-flash',
+      sub_agents=[
+          LlmAgent(name='agent_a', model='gemini-2.5-flash'),
+          LlmAgent(name='agent_b', model='gemini-2.5-flash'),
+      ],
+  )
+
+  with pytest.raises(ValueError) as exc_info:
+    root_agent._LlmAgent__get_agent_to_run('nonexistent_agent')
+
+  error_msg = str(exc_info.value)
+
+  # Verify error message components
+  assert 'nonexistent_agent' in error_msg
+  assert 'Available agents:' in error_msg
+  assert 'agent_a' in error_msg
+  assert 'agent_b' in error_msg
+  assert 'Possible causes:' in error_msg
+  assert 'Suggested fixes:' in error_msg
+
+
+def test_agent_tree_traversal():
+  """Verify agent tree traversal helper works correctly."""
+  root_agent = LlmAgent(
+      name='orchestrator',
+      model='gemini-2.5-flash',
+      sub_agents=[
+          LlmAgent(
+              name='parent_agent',
+              model='gemini-2.5-flash',
+              sub_agents=[
+                  LlmAgent(name='child_agent', model='gemini-2.5-flash'),
+              ],
+          ),
+      ],
+  )
+
+  available_agents = root_agent._get_available_agent_names()
+
+  # Verify all agents in tree are found
+  assert 'orchestrator' in available_agents
+  assert 'parent_agent' in available_agents
+  assert 'child_agent' in available_agents
+  assert len(available_agents) == 3
+
+
+def test_agent_not_found_shows_all_agents():
+  """Verify error message shows all agents (no truncation)."""
+  # Create 100 sub-agents
+  sub_agents = [
+      LlmAgent(name=f'agent_{i}', model='gemini-2.5-flash') for i in range(100)
+  ]
+
+  root_agent = LlmAgent(
+      name='root', model='gemini-2.5-flash', sub_agents=sub_agents
+  )
+
+  with pytest.raises(ValueError) as exc_info:
+    root_agent._LlmAgent__get_agent_to_run('nonexistent')
+
+  error_msg = str(exc_info.value)
+
+  # Verify all agents are shown (no truncation)
+  assert 'agent_0' in error_msg  # First agent shown
+  assert 'agent_99' in error_msg  # Last agent also shown
+  assert 'showing first 20 of' not in error_msg  # No truncation message

tests/unittests/agents/test_llm_agent_fields.py

@@ -0,0 +1,605 @@
+# 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.
+
+"""Unit tests for canonical_xxx fields in LlmAgent."""
+
+import logging
+from typing import Any
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.models.anthropic_llm import Claude
+from google.adk.models.google_llm import Gemini
+from google.adk.models.lite_llm import LiteLlm
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.registry import LLMRegistry
+from google.adk.planners.built_in_planner import BuiltInPlanner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.google_search_tool import google_search
+from google.adk.tools.google_search_tool import GoogleSearchTool
+from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+
+async def _create_readonly_context(
+    agent: LlmAgent, state: Optional[dict[str, Any]] = None
+) -> ReadonlyContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user', state=state
+  )
+  invocation_context = InvocationContext(
+      invocation_id='test_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  return ReadonlyContext(invocation_context)
+
+
+@pytest.mark.parametrize(
+    ('default_model', 'expected_model_name', 'expected_model_type'),
+    [
+        (LlmAgent.DEFAULT_MODEL, LlmAgent.DEFAULT_MODEL, Gemini),
+        ('gemini-2.5-flash', 'gemini-2.5-flash', Gemini),
+    ],
+)
+def test_canonical_model_default_fallback(
+    default_model, expected_model_name, expected_model_type
+):
+  original_default = LlmAgent._default_model
+  LlmAgent.set_default_model(default_model)
+  try:
+    agent = LlmAgent(name='test_agent')
+    assert isinstance(agent.canonical_model, expected_model_type)
+    assert agent.canonical_model.model == expected_model_name
+  finally:
+    LlmAgent.set_default_model(original_default)
+
+
+def test_canonical_model_str():
+  agent = LlmAgent(name='test_agent', model='gemini-pro')
+
+  assert agent.canonical_model.model == 'gemini-pro'
+
+
+def test_canonical_model_llm():
+  llm = LLMRegistry.new_llm('gemini-pro')
+  agent = LlmAgent(name='test_agent', model=llm)
+
+  assert agent.canonical_model == llm
+
+
+def test_canonical_model_inherit():
+  sub_agent = LlmAgent(name='sub_agent')
+  parent_agent = LlmAgent(
+      name='parent_agent', model='gemini-pro', sub_agents=[sub_agent]
+  )
+
+  assert sub_agent.canonical_model == parent_agent.canonical_model
+
+
+def test_canonical_live_model_default_fallback():
+  original_default = LlmAgent._default_live_model
+  LlmAgent.set_default_live_model('gemini-2.0-flash')
+  try:
+    agent = LlmAgent(name='test_agent')
+    assert agent.canonical_live_model.model == 'gemini-2.0-flash'
+  finally:
+    LlmAgent.set_default_live_model(original_default)
+
+
+def test_canonical_live_model_str():
+  agent = LlmAgent(name='test_agent', model='gemini-pro')
+  assert agent.canonical_live_model.model == 'gemini-pro'
+
+
+def test_canonical_live_model_llm():
+  llm = LLMRegistry.new_llm('gemini-pro')
+  agent = LlmAgent(name='test_agent', model=llm)
+  assert agent.canonical_live_model == llm
+
+
+def test_canonical_live_model_inherit():
+  sub_agent = LlmAgent(name='sub_agent')
+  parent_agent = LlmAgent(
+      name='parent_agent', model='gemini-pro', sub_agents=[sub_agent]
+  )
+  assert sub_agent.canonical_live_model == parent_agent.canonical_live_model
+
+
+async def test_canonical_instruction_str():
+  agent = LlmAgent(name='test_agent', instruction='instruction')
+  ctx = await _create_readonly_context(agent)
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_instruction(ctx)
+  )
+  assert canonical_instruction == 'instruction'
+  assert not bypass_state_injection
+
+
+async def test_canonical_instruction():
+  def _instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_instruction(ctx)
+  )
+  assert canonical_instruction == 'instruction: state_value'
+  assert bypass_state_injection
+
+
+async def test_async_canonical_instruction():
+  async def _instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_instruction(ctx)
+  )
+  assert canonical_instruction == 'instruction: state_value'
+  assert bypass_state_injection
+
+
+async def test_canonical_global_instruction_str():
+  agent = LlmAgent(name='test_agent', global_instruction='global instruction')
+  ctx = await _create_readonly_context(agent)
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_global_instruction(ctx)
+  )
+  assert canonical_instruction == 'global instruction'
+  assert not bypass_state_injection
+
+
+async def test_canonical_global_instruction():
+  def _global_instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'global instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(
+      name='test_agent', global_instruction=_global_instruction_provider
+  )
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+
+  canonical_global_instruction, bypass_state_injection = (
+      await agent.canonical_global_instruction(ctx)
+  )
+  assert canonical_global_instruction == 'global instruction: state_value'
+  assert bypass_state_injection
+
+
+async def test_async_canonical_global_instruction():
+  async def _global_instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'global instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(
+      name='test_agent', global_instruction=_global_instruction_provider
+  )
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+  canonical_global_instruction, bypass_state_injection = (
+      await agent.canonical_global_instruction(ctx)
+  )
+  assert canonical_global_instruction == 'global instruction: state_value'
+  assert bypass_state_injection
+
+
+def test_output_schema_with_sub_agents_will_not_throw():
+  class Schema(BaseModel):
+    pass
+
+  sub_agent = LlmAgent(
+      name='sub_agent',
+  )
+
+  agent = LlmAgent(
+      name='test_agent',
+      output_schema=Schema,
+      sub_agents=[sub_agent],
+  )
+
+  # Transfer is not disabled
+  assert not agent.disallow_transfer_to_parent
+  assert not agent.disallow_transfer_to_peers
+
+  assert agent.output_schema == Schema
+  assert agent.sub_agents == [sub_agent]
+
+
+def test_output_schema_with_tools_will_not_throw():
+  class Schema(BaseModel):
+    pass
+
+  def _a_tool():
+    pass
+
+  LlmAgent(
+      name='test_agent',
+      output_schema=Schema,
+      tools=[_a_tool],
+  )
+
+
+def test_before_model_callback():
+  def _before_model_callback(
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> None:
+    return None
+
+  agent = LlmAgent(
+      name='test_agent', before_model_callback=_before_model_callback
+  )
+
+  # TODO: add more logic assertions later.
+  assert agent.before_model_callback is not None
+
+
+def test_validate_generate_content_config_thinking_config_allow():
+  """Tests that thinking_config is now allowed directly in the agent init."""
+  agent = LlmAgent(
+      name='test_agent',
+      generate_content_config=types.GenerateContentConfig(
+          thinking_config=types.ThinkingConfig(include_thoughts=True)
+      ),
+  )
+  assert agent.generate_content_config.thinking_config.include_thoughts is True
+
+
+def test_thinking_config_precedence_warning():
+  """Tests that a UserWarning is issued when both manual config and planner exist."""
+
+  config = types.GenerateContentConfig(
+      thinking_config=types.ThinkingConfig(include_thoughts=True)
+  )
+  planner = BuiltInPlanner(
+      thinking_config=types.ThinkingConfig(include_thoughts=True)
+  )
+
+  with pytest.warns(
+      UserWarning, match="planner's configuration will take precedence"
+  ):
+    LlmAgent(name='test_agent', generate_content_config=config, planner=planner)
+
+
+def test_validate_generate_content_config_tools_throw():
+  """Tests that tools cannot be set directly in config."""
+  with pytest.raises(ValueError):
+    _ = LlmAgent(
+        name='test_agent',
+        generate_content_config=types.GenerateContentConfig(
+            tools=[types.Tool(function_declarations=[])]
+        ),
+    )
+
+
+def test_validate_generate_content_config_system_instruction_throw():
+  """Tests that system instructions cannot be set directly in config."""
+  with pytest.raises(ValueError):
+    _ = LlmAgent(
+        name='test_agent',
+        generate_content_config=types.GenerateContentConfig(
+            system_instruction='system instruction'
+        ),
+    )
+
+
+def test_validate_generate_content_config_response_schema_throw():
+  """Tests that response schema cannot be set directly in config."""
+
+  class Schema(BaseModel):
+    pass
+
+  with pytest.raises(ValueError):
+    _ = LlmAgent(
+        name='test_agent',
+        generate_content_config=types.GenerateContentConfig(
+            response_schema=Schema
+        ),
+    )
+
+
+def test_allow_transfer_by_default():
+  sub_agent = LlmAgent(name='sub_agent')
+  agent = LlmAgent(name='test_agent', sub_agents=[sub_agent])
+
+  assert not agent.disallow_transfer_to_parent
+  assert not agent.disallow_transfer_to_peers
+
+
+# TODO(b/448114567): Remove TestCanonicalTools once the workaround
+# is no longer needed.
+class TestCanonicalTools:
+  """Unit tests for canonical_tools in LlmAgent."""
+
+  @staticmethod
+  def _my_tool(sides: int) -> int:
+    return sides
+
+  async def test_handle_google_search_with_other_tools(self):
+    """Test that google_search is wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            GoogleSearchTool(bypass_multi_tools_limit=True),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'google_search_agent'
+    assert tools[1].__class__.__name__ == 'GoogleSearchAgentTool'
+
+  async def test_handle_google_search_with_other_tools_no_bypass(self):
+    """Test that google_search is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            GoogleSearchTool(bypass_multi_tools_limit=False),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'google_search'
+    assert tools[1].__class__.__name__ == 'GoogleSearchTool'
+
+  async def test_handle_google_search_only(self):
+    """Test that google_search is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            google_search,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == 'google_search'
+    assert tools[0].__class__.__name__ == 'GoogleSearchTool'
+
+  async def test_function_tool_only(self):
+    """Test that function tool is not affected."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+
+  @mock.patch(
+      'google.auth.default',
+      mock.MagicMock(return_value=('credentials', 'project')),
+  )
+  async def test_handle_vais_with_other_tools(self):
+    """Test that VertexAiSearchTool is replaced with Discovery Engine Search."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            VertexAiSearchTool(
+                data_store_id='test_data_store_id',
+                bypass_multi_tools_limit=True,
+            ),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'discovery_engine_search'
+    assert tools[1].__class__.__name__ == 'DiscoveryEngineSearchTool'
+
+  async def test_handle_vais_with_other_tools_no_bypass(self):
+    """Test that VertexAiSearchTool is not replaced."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            VertexAiSearchTool(
+                data_store_id='test_data_store_id',
+                bypass_multi_tools_limit=False,
+            ),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'vertex_ai_search'
+    assert tools[1].__class__.__name__ == 'VertexAiSearchTool'
+
+  async def test_handle_vais_only(self):
+    """Test that VertexAiSearchTool is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            VertexAiSearchTool(data_store_id='test_data_store_id'),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == 'vertex_ai_search'
+    assert tools[0].__class__.__name__ == 'VertexAiSearchTool'
+
+  async def test_multiple_tools_resolution(self):
+    """Test that multiple tools are resolved correctly."""
+
+    def _tool_1():
+      pass
+
+    def _tool_2():
+      pass
+
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[_tool_1, _tool_2],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_tool_1'
+    assert tools[1].name == '_tool_2'
+
+  async def test_canonical_tools_graceful_degradation_on_toolset_error(self):
+    """Test that canonical_tools returns tools from working toolsets when one fails."""
+    from google.adk.tools.base_tool import BaseTool
+    from google.adk.tools.base_toolset import BaseToolset
+
+    class FailingToolset(BaseToolset):
+
+      async def get_tools(self, readonly_context=None):
+        raise ConnectionError('MCP server unavailable')
+
+    class WorkingToolset(BaseToolset):
+
+      async def get_tools(self, readonly_context=None):
+        tool = mock.MagicMock(spec=BaseTool)
+        tool.name = 'working_tool'
+        tool._get_declaration = mock.MagicMock(return_value=None)
+        return [tool]
+
+    def _regular_tool():
+      pass
+
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[_regular_tool, FailingToolset(), WorkingToolset()],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    # Should have the regular tool + working toolset tool, but not crash
+    assert len(tools) == 2
+    assert tools[0].name == '_regular_tool'
+    assert tools[1].name == 'working_tool'
+
+
+# Tests for multi-provider model support via string model names
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'gemini-2.5-flash',
+        'gemini-2.5-pro',
+    ],
+)
+def test_agent_with_gemini_string_model(model_name):
+  """Test that Agent accepts Gemini model strings and resolves to Gemini."""
+  agent = LlmAgent(name='test_agent', model=model_name)
+  assert isinstance(agent.canonical_model, Gemini)
+  assert agent.canonical_model.model == model_name
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'claude-3-5-sonnet-v2@20241022',
+        'claude-sonnet-4@20250514',
+    ],
+)
+def test_agent_with_claude_string_model(model_name):
+  """Test that Agent accepts Claude model strings and resolves to Claude."""
+  agent = LlmAgent(name='test_agent', model=model_name)
+  assert isinstance(agent.canonical_model, Claude)
+  assert agent.canonical_model.model == model_name
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'openai/gpt-4o',
+        'groq/llama3-70b-8192',
+        'anthropic/claude-3-opus-20240229',
+    ],
+)
+def test_agent_with_litellm_string_model(model_name):
+  """Test that Agent accepts LiteLLM provider strings."""
+  agent = LlmAgent(name='test_agent', model=model_name)
+  assert isinstance(agent.canonical_model, LiteLlm)
+  assert agent.canonical_model.model == model_name
+
+
+def test_builtin_planner_overwrite_logging(caplog):
+  """Tests that the planner logs an DEBUG message when overwriting a config."""
+
+  planner = BuiltInPlanner(
+      thinking_config=types.ThinkingConfig(include_thoughts=True)
+  )
+
+  # Create a request that already has a thinking_config
+  req = LlmRequest(
+      contents=[],
+      config=types.GenerateContentConfig(
+          thinking_config=types.ThinkingConfig(include_thoughts=True)
+      ),
+  )
+
+  with caplog.at_level(
+      logging.DEBUG, logger='google_adk.google.adk.planners.built_in_planner'
+  ):
+    planner.apply_thinking_config(req)
+  assert (
+      'Overwriting `thinking_config` from `generate_content_config`'
+      in caplog.text
+  )

tests/unittests/agents/test_llm_agent_include_contents.py

@@ -0,0 +1,434 @@
+# 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.
+
+"""Unit tests for LlmAgent include_contents field behavior."""
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.genai import types
+import pytest
+
+from .. import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_include_contents_default_behavior():
+  """Test that include_contents='default' preserves conversation history including tool interactions."""
+
+  def simple_tool(message: str) -> dict:
+    return {"result": f"Tool processed: {message}"}
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "first"}
+          ),
+          "First response",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "second"}
+          ),
+          "Second response",
+      ]
+  )
+
+  agent = LlmAgent(
+      name="test_agent",
+      model=mock_model,
+      include_contents="default",
+      instruction="You are a helpful assistant",
+      tools=[simple_tool],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  runner.run("First message")
+  runner.run("Second message")
+
+  # First turn requests
+  assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
+      ("user", "First message")
+  ]
+
+  assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
+      ("user", "First message"),
+      (
+          "model",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "first"}
+          ),
+      ),
+      (
+          "user",
+          types.Part.from_function_response(
+              name="simple_tool", response={"result": "Tool processed: first"}
+          ),
+      ),
+  ]
+
+  # Second turn should include full conversation history
+  assert testing_utils.simplify_contents(mock_model.requests[2].contents) == [
+      ("user", "First message"),
+      (
+          "model",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "first"}
+          ),
+      ),
+      (
+          "user",
+          types.Part.from_function_response(
+              name="simple_tool", response={"result": "Tool processed: first"}
+          ),
+      ),
+      ("model", "First response"),
+      ("user", "Second message"),
+  ]
+
+  # Second turn with tool should include full history + current tool interaction
+  assert testing_utils.simplify_contents(mock_model.requests[3].contents) == [
+      ("user", "First message"),
+      (
+          "model",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "first"}
+          ),
+      ),
+      (
+          "user",
+          types.Part.from_function_response(
+              name="simple_tool", response={"result": "Tool processed: first"}
+          ),
+      ),
+      ("model", "First response"),
+      ("user", "Second message"),
+      (
+          "model",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "second"}
+          ),
+      ),
+      (
+          "user",
+          types.Part.from_function_response(
+              name="simple_tool", response={"result": "Tool processed: second"}
+          ),
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_include_contents_none_behavior():
+  """Test that include_contents='none' excludes conversation history but includes current input."""
+
+  def simple_tool(message: str) -> dict:
+    return {"result": f"Tool processed: {message}"}
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "first"}
+          ),
+          "First response",
+          "Second response",
+      ]
+  )
+
+  agent = LlmAgent(
+      name="test_agent",
+      model=mock_model,
+      include_contents="none",
+      instruction="You are a helpful assistant",
+      tools=[simple_tool],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  runner.run("First message")
+  runner.run("Second message")
+
+  # First turn behavior
+  assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
+      ("user", "First message")
+  ]
+
+  assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
+      ("user", "First message"),
+      (
+          "model",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "first"}
+          ),
+      ),
+      (
+          "user",
+          types.Part.from_function_response(
+              name="simple_tool", response={"result": "Tool processed: first"}
+          ),
+      ),
+  ]
+
+  # Second turn should only have current input, no history
+  assert testing_utils.simplify_contents(mock_model.requests[2].contents) == [
+      ("user", "Second message")
+  ]
+
+  # System instruction and tools should be preserved
+  assert (
+      "You are a helpful assistant"
+      in mock_model.requests[0].config.system_instruction
+  )
+  assert len(mock_model.requests[0].config.tools) > 0
+
+
+def test_model_input_context_is_sent_to_model_without_persisting_to_session():
+  mock_model = testing_utils.MockModel.create(responses=["Answer"])
+  agent = LlmAgent(name="test_agent", model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent)
+  session = runner.session
+
+  list(
+      runner.runner.run(
+          user_id=session.user_id,
+          session_id=session.id,
+          new_message=testing_utils.get_user_content("Question"),
+          run_config=RunConfig(
+              model_input_context=[
+                  types.UserContent("Relevant context for this turn")
+              ]
+          ),
+      )
+  )
+
+  assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
+      ("user", "Relevant context for this turn"),
+      ("user", "Question"),
+  ]
+  assert testing_utils.simplify_events(runner.session.events) == [
+      ("user", "Question"),
+      ("test_agent", "Answer"),
+  ]
+
+
+def test_model_input_context_stays_before_user_message_after_tool_call():
+  def simple_tool(message: str) -> dict:
+    return {"result": f"Tool processed: {message}"}
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "payload"}
+          ),
+          "Answer",
+      ]
+  )
+  agent = LlmAgent(name="test_agent", model=mock_model, tools=[simple_tool])
+  runner = testing_utils.InMemoryRunner(agent)
+  session = runner.session
+
+  list(
+      runner.runner.run(
+          user_id=session.user_id,
+          session_id=session.id,
+          new_message=testing_utils.get_user_content("Question"),
+          run_config=RunConfig(
+              model_input_context=[
+                  types.UserContent("Relevant context for this turn")
+              ]
+          ),
+      )
+  )
+
+  assert testing_utils.simplify_contents(mock_model.requests[0].contents) == [
+      ("user", "Relevant context for this turn"),
+      ("user", "Question"),
+  ]
+  assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
+      ("user", "Relevant context for this turn"),
+      ("user", "Question"),
+      (
+          "model",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "payload"}
+          ),
+      ),
+      (
+          "user",
+          types.Part.from_function_response(
+              name="simple_tool",
+              response={"result": "Tool processed: payload"},
+          ),
+      ),
+  ]
+  assert testing_utils.simplify_events(runner.session.events) == [
+      ("user", "Question"),
+      (
+          "test_agent",
+          types.Part.from_function_call(
+              name="simple_tool", args={"message": "payload"}
+          ),
+      ),
+      (
+          "test_agent",
+          types.Part.from_function_response(
+              name="simple_tool",
+              response={"result": "Tool processed: payload"},
+          ),
+      ),
+      ("test_agent", "Answer"),
+  ]
+
+
+def test_model_input_context_with_include_contents_none_sub_agent():
+  agent1_model = testing_utils.MockModel.create(
+      responses=["Agent1 response: XYZ"]
+  )
+  agent1 = LlmAgent(name="agent1", model=agent1_model)
+
+  agent2_model = testing_utils.MockModel.create(
+      responses=["Agent2 final response"]
+  )
+  agent2 = LlmAgent(
+      name="agent2",
+      model=agent2_model,
+      include_contents="none",
+  )
+  sequential_agent = SequentialAgent(
+      name="sequential_test_agent", sub_agents=[agent1, agent2]
+  )
+  runner = testing_utils.InMemoryRunner(sequential_agent)
+  session = runner.session
+
+  list(
+      runner.runner.run(
+          user_id=session.user_id,
+          session_id=session.id,
+          new_message=testing_utils.get_user_content("Original user request"),
+          run_config=RunConfig(
+              model_input_context=[
+                  types.UserContent("Relevant context for this turn")
+              ]
+          ),
+      )
+  )
+
+  assert testing_utils.simplify_contents(agent1_model.requests[0].contents) == [
+      ("user", "Relevant context for this turn"),
+      ("user", "Original user request"),
+  ]
+  assert testing_utils.simplify_contents(agent2_model.requests[0].contents) == [
+      ("user", "Relevant context for this turn"),
+      (
+          "user",
+          [
+              types.Part(text="For context:"),
+              types.Part(text="[agent1] said: Agent1 response: XYZ"),
+          ],
+      ),
+  ]
+
+
+def test_model_input_context_without_user_message_is_prepended_before_history():
+  mock_model = testing_utils.MockModel.create(
+      responses=["First answer", "Second answer"]
+  )
+  agent = LlmAgent(name="test_agent", model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent)
+  session = runner.session
+
+  list(
+      runner.runner.run(
+          user_id=session.user_id,
+          session_id=session.id,
+          new_message=testing_utils.get_user_content("First question"),
+      )
+  )
+  # No new_message, so the invocation has no user content to anchor before
+  # (e.g. live mode or a re-run over existing history). The transient context
+  # falls back to the front of the request, before all prior history.
+  list(
+      runner.runner.run(
+          user_id=session.user_id,
+          session_id=session.id,
+          new_message=None,
+          run_config=RunConfig(
+              model_input_context=[
+                  types.UserContent("Relevant context for this turn")
+              ]
+          ),
+      )
+  )
+
+  assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [
+      ("user", "Relevant context for this turn"),
+      ("user", "First question"),
+      ("model", "First answer"),
+  ]
+  assert testing_utils.simplify_events(runner.session.events) == [
+      ("user", "First question"),
+      ("test_agent", "First answer"),
+      ("test_agent", "Second answer"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_include_contents_none_sequential_agents():
+  """Test include_contents='none' with sequential agents."""
+
+  agent1_model = testing_utils.MockModel.create(
+      responses=["Agent1 response: XYZ"]
+  )
+  agent1 = LlmAgent(
+      name="agent1",
+      model=agent1_model,
+      instruction="You are Agent1",
+  )
+
+  agent2_model = testing_utils.MockModel.create(
+      responses=["Agent2 final response"]
+  )
+  agent2 = LlmAgent(
+      name="agent2",
+      model=agent2_model,
+      include_contents="none",
+      instruction="You are Agent2",
+  )
+
+  sequential_agent = SequentialAgent(
+      name="sequential_test_agent", sub_agents=[agent1, agent2]
+  )
+
+  runner = testing_utils.InMemoryRunner(sequential_agent)
+  events = runner.run("Original user request")
+
+  simplified_events = [event for event in events if event.content]
+  assert len(simplified_events) == 2
+  assert "Agent1 response" in str(simplified_events[0].content)
+  assert "Agent2 final response" in str(simplified_events[1].content)
+
+  # Agent1 sees original user request
+  agent1_contents = testing_utils.simplify_contents(
+      agent1_model.requests[0].contents
+  )
+  assert ("user", "Original user request") in agent1_contents
+
+  # Agent2 with include_contents='none' should not see original request
+  agent2_contents = testing_utils.simplify_contents(
+      agent2_model.requests[0].contents
+  )
+
+  assert not any(
+      "Original user request" in str(content) for _, content in agent2_contents
+  )
+  assert any(
+      "Agent1 response" in str(content) for _, content in agent2_contents
+  )

tests/unittests/agents/test_llm_agent_interruptions.py

@@ -0,0 +1,481 @@
+# 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.
+
+from __future__ import annotations
+
+from google.adk.agents import LlmAgent
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.long_running_tool import LongRunningFunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from tests.unittests import testing_utils
+from tests.unittests.agents.llm.event_utils import text_parts
+
+_USER_ID = 'test_user'
+_SESSION_ID = 'test_session'
+
+
+async def _setup_runner(mock_model, tools=None, **agent_kwargs):
+  """Setup runner with LlmAgent directly."""
+  llm_agent = LlmAgent(
+      name='test_agent',
+      model=mock_model,
+      tools=tools or [],
+      **agent_kwargs,
+  )
+  session_service = InMemorySessionService()
+  await session_service.create_session(
+      app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+  )
+  runner = Runner(
+      app_name='test',
+      agent=llm_agent,
+      session_service=session_service,
+  )
+  return runner
+
+
+async def _run_turn(runner, user_message):
+  """Run a single turn."""
+  return [
+      e
+      async for e in runner.run_async(
+          user_id=_USER_ID,
+          session_id=_SESSION_ID,
+          new_message=types.Content(
+              role='user', parts=[types.Part(text=user_message)]
+          ),
+      )
+  ]
+
+
+async def _resume_turn(
+    runner, prev_events, tool_name, tool_response_value='done'
+):
+  """Resume after an interrupt."""
+  fc_ids = []
+  for e in prev_events:
+    if e.content and e.content.parts:
+      for p in e.content.parts:
+        if (
+            p.function_call
+            and p.function_call.name == tool_name
+            and p.function_call.id
+        ):
+          fc_ids.append(p.function_call.id)
+    if getattr(e.output, 'function_calls', None):
+      for fc in e.output.function_calls:
+        if fc.name == tool_name and fc.id:
+          fc_ids.append(fc.id)
+
+  if not fc_ids:
+    for e in prev_events:
+      if e.long_running_tool_ids:
+        fc_ids = list(e.long_running_tool_ids)
+        break
+
+  invocation_id = prev_events[0].invocation_id
+
+  fr_parts = [
+      types.Part(
+          function_response=types.FunctionResponse(
+              name=tool_name,
+              id=fc_id,
+              response={'result': tool_response_value},
+          )
+      )
+      for fc_id in fc_ids
+  ]
+  resume_msg = types.Content(role='user', parts=fr_parts)
+
+  return [
+      e
+      async for e in runner.run_async(
+          user_id=_USER_ID,
+          session_id=_SESSION_ID,
+          invocation_id=invocation_id,
+          new_message=resume_msg,
+      )
+  ]
+
+
+def create_lro_tool(name: str = 'long_running_op') -> LongRunningFunctionTool:
+  """Creates a minimal LRO tool for testing."""
+
+  def _impl() -> None:
+    return None
+
+  _impl.__name__ = name
+  return LongRunningFunctionTool(_impl)
+
+
+# ---------------------------------------------------------------------------
+# Tests: Single Agent
+# ---------------------------------------------------------------------------
+
+
+class TestSingleAgentInterruptions:
+  """Tests for single agent triggering interruptions."""
+
+  async def test_single_agent_yields_on_long_running_tool(self):
+    """Single agent yields on Long Running Tool.
+
+    Arrange: Set up a single agent with a long running tool.
+    Act: Run the agent with a prompt that triggers the tool.
+    Assert: Verify that the execution yields a long running tool interrupt.
+    """
+
+    fc = types.Part.from_function_call(name='long_running_op', args={})
+    mock_model = testing_utils.MockModel.create(responses=[fc, 'Final answer'])
+
+    lro_tool = create_lro_tool()
+    runner = await _setup_runner(mock_model, tools=[lro_tool])
+
+    # Act: Run first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Assert: Should have triggered function call
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'long_running_op'
+            for p in e.content.parts or []
+        )
+        for e in events
+    )
+    assert len(mock_model.requests) == 1
+
+    # Act: Resume
+    resume_events = await _resume_turn(runner, events, 'long_running_op')
+
+    # Assert: Should have completed
+    assert any('Final answer' in t for t in text_parts(resume_events))
+    assert len(mock_model.requests) == 2
+
+  async def test_single_agent_request_input_tool_interrupt_and_resume(self):
+    """Test that using RequestInputTool successfully triggers an interrupt and resumes with user input."""
+    from google.adk.tools import request_input
+
+    fc = types.Part.from_function_call(
+        name='adk_request_input',
+        args={'message': 'Which file?', 'response_schema': {'type': 'string'}},
+    )
+    mock_model = testing_utils.MockModel.create(
+        responses=[fc, 'Continuing with file: file_a.txt']
+    )
+
+    runner = await _setup_runner(mock_model, tools=[request_input])
+
+    # Act: Run first turn
+    events = await _run_turn(runner, 'Start')
+
+    # Assert: Verify the interrupt event is produced
+    assert any(e.long_running_tool_ids for e in events)
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in e.content.parts or []
+        )
+        for e in events
+    )
+
+    # Act: Resume with the response
+    resume_events = await _resume_turn(
+        runner, events, 'adk_request_input', tool_response_value='file_a.txt'
+    )
+
+    # Assert: Execution should continue with user response in the prompt history
+    assert len(mock_model.requests) == 2
+
+    # Assert: Verify the second request contains the FunctionCall & FunctionResponse pair
+    second_req_contents = mock_model.requests[1].contents
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+    assert any(
+        any(
+            p.function_response
+            and p.function_response.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+
+  async def test_single_agent_request_input_tool_structured_schema(self):
+    """Test that using RequestInputTool with a structured object schema successfully interrupts and resumes with a dictionary response."""
+    from google.adk.tools import request_input
+
+    schema = {
+        'type': 'object',
+        'properties': {
+            'host': {'type': 'string'},
+            'port': {'type': 'integer'},
+        },
+        'required': ['host'],
+    }
+    fc = types.Part.from_function_call(
+        name='adk_request_input',
+        args={
+            'message': 'Provide DB connection details:',
+            'response_schema': schema,
+        },
+    )
+    mock_model = testing_utils.MockModel.create(
+        responses=[fc, 'Connected to localhost:3306']
+    )
+
+    runner = await _setup_runner(mock_model, tools=[request_input])
+
+    # Act: Run first turn
+    events = await _run_turn(runner, 'Start')
+
+    # Assert: Verify the interrupt event is produced with the schema args
+    assert any(e.long_running_tool_ids for e in events)
+    fc_event = next(
+        e
+        for e in events
+        if e.content
+        and any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in e.content.parts or []
+        )
+    )
+    fc_part = next(p for p in fc_event.content.parts if p.function_call)
+    assert fc_part.function_call.args['response_schema'] == schema
+
+    # Act: Resume with a structured dict response
+    db_details = {'host': 'localhost', 'port': 3306}
+    resume_events = await _resume_turn(
+        runner, events, 'adk_request_input', tool_response_value=db_details
+    )
+
+    # Assert: Execution should continue with the structured user response
+    assert len(mock_model.requests) == 2
+
+    # Assert: Verify the second request contains the FunctionCall & FunctionResponse pair
+    second_req_contents = mock_model.requests[1].contents
+    assert any(
+        any(
+            p.function_call and p.function_call.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+    assert any(
+        any(
+            p.function_response
+            and p.function_response.name == 'adk_request_input'
+            for p in c.parts or []
+        )
+        for c in second_req_contents
+    )
+
+
+class TestNestedAgentInterruptions:
+  """Tests for multi-agent setups with interruptions."""
+
+  async def test_child_agent_interrupt_and_resume(self):
+    """Child agent yields on LRO and resumes successfully.
+
+    Arrange: Parent agent with Child agent. Parent transfers to Child.
+      Child calls LRO tool.
+    Act: Run, expect LRO interrupt. Then resume.
+    Assert: Should complete successfully.
+    """
+
+    def transfer_to_child(tool_context: ToolContext) -> str:
+      tool_context.actions.transfer_to_agent = 'child_agent'
+      return 'transferring'
+
+    # Child agent
+    fc_child = types.Part.from_function_call(name='child_lro', args={})
+    child_mock_model = testing_utils.MockModel.create(
+        responses=[fc_child, 'Child final answer']
+    )
+
+    lro_tool = create_lro_tool('child_lro')
+
+    child_agent = LlmAgent(
+        name='child_agent',
+        model=child_mock_model,
+        tools=[lro_tool],
+    )
+
+    # Parent agent
+    fc_parent = types.Part.from_function_call(name='transfer_to_child', args={})
+    parent_mock_model = testing_utils.MockModel.create(
+        responses=[fc_parent, fc_parent, 'Parent final answer']
+    )
+
+    parent_agent = LlmAgent(
+        name='parent_agent',
+        model=parent_mock_model,
+        tools=[transfer_to_child],
+        sub_agents=[child_agent],
+    )
+
+    # Setup runner
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+    )
+    runner = Runner(
+        app_name='test', agent=parent_agent, session_service=session_service
+    )
+
+    # When Parent runs the first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Then it should trigger child LRO interrupt
+    assert any(e.long_running_tool_ids for e in events)
+
+    # When Parent resumes the turn
+    resume_events = await _resume_turn(runner, events, 'child_lro')
+
+    # Then it should complete successfully
+    assert any('Child final answer' in t for t in text_parts(resume_events))
+
+  @pytest.mark.xfail(reason='Task agent as subagent not supported yet.')
+  async def test_task_child_agent_interrupt_and_resume(self):
+    """Task child agent yields on LRO and resumes successfully.
+
+    Arrange: Parent agent with Task Child agent. Parent transfers to Child.
+      Child calls LRO tool.
+    Act: Run, expect LRO interrupt. Then resume.
+    Assert: Should complete successfully.
+    """
+
+    def transfer_to_child(tool_context: ToolContext) -> str:
+      tool_context.actions.transfer_to_agent = 'child_agent'
+      return 'transferring'
+
+    # Child agent (Task mode)
+    fc_child = types.Part.from_function_call(name='child_lro', args={})
+    fc_finish = types.Part.from_function_call(
+        name='finish_task', args={'result': 'Task done'}
+    )
+    child_mock_model = testing_utils.MockModel.create(
+        responses=[fc_child, fc_finish, 'Child final answer']
+    )
+
+    lro_tool = create_lro_tool('child_lro')
+
+    child_agent = LlmAgent(
+        name='child_agent',
+        model=child_mock_model,
+        tools=[lro_tool],
+        mode='task',
+    )
+
+    # Parent agent
+    fc_parent = types.Part.from_function_call(name='transfer_to_child', args={})
+    parent_mock_model = testing_utils.MockModel.create(
+        responses=[fc_parent, 'Parent final answer']
+    )
+
+    parent_agent = LlmAgent(
+        name='parent_agent',
+        model=parent_mock_model,
+        tools=[transfer_to_child],
+        sub_agents=[child_agent],
+    )
+
+    # Setup runner
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+    )
+    runner = Runner(
+        app_name='test', agent=parent_agent, session_service=session_service
+    )
+
+    # When Parent runs the first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Then it should trigger child LRO interrupt
+    assert any(e.long_running_tool_ids for e in events)
+
+    # When Parent resumes the turn
+    resume_events = await _resume_turn(runner, events, 'child_lro')
+
+    # Then it should complete successfully
+    assert any('Parent final answer' in t for t in text_parts(resume_events))
+
+  @pytest.mark.xfail(reason='Single-turn agent as subagent not supported yet.')
+  async def test_single_turn_child_agent_interrupt_and_resume(self):
+    """Single-turn child agent yields on LRO and resumes successfully.
+
+    Arrange: Parent agent with Single-turn Child agent.
+      Parent calls Child via tool.
+      Child calls LRO tool.
+    Act: Run, expect LRO interrupt. Then resume.
+    Assert: Should complete successfully.
+    """
+
+    # Child agent (Single-turn)
+    fc_child = types.Part.from_function_call(name='child_lro', args={})
+    child_mock_model = testing_utils.MockModel.create(
+        responses=[fc_child, 'Child final answer']
+    )
+
+    lro_tool = create_lro_tool('child_lro')
+
+    child_agent = LlmAgent(
+        name='child_agent',
+        model=child_mock_model,
+        tools=[lro_tool],
+        mode='single_turn',
+    )
+
+    # Parent agent
+    fc_call_child = types.Part.from_function_call(
+        name='child_agent', args={'request': 'Go to child'}
+    )
+    parent_mock_model = testing_utils.MockModel.create(
+        responses=[fc_call_child, 'Parent final answer']
+    )
+
+    parent_agent = LlmAgent(
+        name='parent_agent',
+        model=parent_mock_model,
+        sub_agents=[child_agent],
+    )
+
+    # Setup runner
+    session_service = InMemorySessionService()
+    await session_service.create_session(
+        app_name='test', user_id=_USER_ID, session_id=_SESSION_ID
+    )
+    runner = Runner(
+        app_name='test', agent=parent_agent, session_service=session_service
+    )
+
+    # When Parent runs the first turn
+    events = await _run_turn(runner, 'Go')
+
+    # Then it should trigger child LRO interrupt
+    assert any(e.long_running_tool_ids for e in events)
+
+    # When Parent resumes the turn
+    resume_events = await _resume_turn(runner, events, 'child_lro')
+
+    # Then it should complete successfully
+    assert any('Parent final answer' in t for t in text_parts(resume_events))

tests/unittests/agents/test_llm_agent_output_save.py

@@ -0,0 +1,357 @@
+# 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.
+
+"""Unit tests for LlmAgent output saving functionality."""
+
+import logging
+from unittest.mock import patch
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.events.event import Event
+from google.adk.events.event_actions import EventActions
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from .. import testing_utils
+
+
+class MockOutputSchema(BaseModel):
+  message: str
+  confidence: float
+
+
+def create_test_event(
+    author: str = "test_agent",
+    content_text: str = "Hello world",
+    is_final: bool = True,
+    invocation_id: str = "test_invocation",
+) -> Event:
+  """Helper to create test events."""
+  # Create mock content
+  parts = [types.Part.from_text(text=content_text)] if content_text else []
+  content = types.Content(role="model", parts=parts) if parts else None
+
+  # Create event
+  event = Event(
+      invocation_id=invocation_id,
+      author=author,
+      content=content,
+      actions=EventActions(),
+  )
+
+  # Mock is_final_response if needed
+  if not is_final:
+    event.partial = True
+
+  return event
+
+
+class TestLlmAgentOutputSave:
+  """Test suite for LlmAgent output saving functionality."""
+
+  def test_maybe_save_output_to_state_skips_different_author(self, caplog):
+    """Test that output is not saved when event author differs from agent name."""
+    # Set the LlmAgent logger to DEBUG level
+    llm_agent_logger = logging.getLogger(
+        "google_adk.google.adk.agents.llm_agent"
+    )
+    original_level = llm_agent_logger.level
+    llm_agent_logger.setLevel(logging.DEBUG)
+
+    try:
+      agent = LlmAgent(name="agent_a", output_key="result")
+      event = create_test_event(
+          author="agent_b", content_text="Response from B"
+      )
+
+      with caplog.at_level("DEBUG"):
+        agent._LlmAgent__maybe_save_output_to_state(event)
+
+      # Should not add anything to state_delta
+      assert len(event.actions.state_delta) == 0
+
+      # Should log the skip
+      assert (
+          "Skipping output save for agent agent_a: event authored by agent_b"
+          in caplog.text
+      )
+    finally:
+      # Restore original logger level
+      llm_agent_logger.setLevel(original_level)
+
+  def test_maybe_save_output_to_state_saves_same_author(self):
+    """Test that output is saved when event author matches agent name."""
+    agent = LlmAgent(name="test_agent", output_key="result")
+    event = create_test_event(author="test_agent", content_text="Test response")
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should save to state_delta
+    assert event.actions.state_delta["result"] == "Test response"
+
+  def test_maybe_save_output_to_state_no_output_key(self):
+    """Test that nothing is saved when output_key is not set."""
+    agent = LlmAgent(name="test_agent")  # No output_key
+    event = create_test_event(author="test_agent", content_text="Test response")
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should not save anything
+    assert len(event.actions.state_delta) == 0
+
+  def test_maybe_save_output_to_state_not_final_response(self):
+    """Test that output is not saved for non-final responses."""
+    agent = LlmAgent(name="test_agent", output_key="result")
+    event = create_test_event(
+        author="test_agent", content_text="Partial response", is_final=False
+    )
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should not save partial responses
+    assert len(event.actions.state_delta) == 0
+
+  def test_maybe_save_output_to_state_no_content(self):
+    """Test that nothing is saved when event has no content."""
+    agent = LlmAgent(name="test_agent", output_key="result")
+    event = create_test_event(author="test_agent", content_text="")
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should not save empty content
+    assert len(event.actions.state_delta) == 0
+
+  def test_maybe_save_output_to_state_with_output_schema(self):
+    """Test that output is processed with schema when output_schema is set."""
+    agent = LlmAgent(
+        name="test_agent", output_key="result", output_schema=MockOutputSchema
+    )
+
+    # Create event with JSON content
+    json_content = '{"message": "Hello", "confidence": 0.95}'
+    event = create_test_event(author="test_agent", content_text=json_content)
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should save parsed and validated output
+    expected_output = {"message": "Hello", "confidence": 0.95}
+    assert event.actions.state_delta["result"] == expected_output
+
+  def test_maybe_save_output_to_state_multiple_parts(self):
+    """Test that multiple text parts are concatenated."""
+    agent = LlmAgent(name="test_agent", output_key="result")
+
+    # Create event with multiple text parts
+    parts = [
+        types.Part.from_text(text="Hello "),
+        types.Part.from_text(text="world"),
+        types.Part.from_text(text="!"),
+    ]
+    content = types.Content(role="model", parts=parts)
+
+    event = Event(
+        invocation_id="test_invocation",
+        author="test_agent",
+        content=content,
+        actions=EventActions(),
+    )
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should concatenate all text parts
+    assert event.actions.state_delta["result"] == "Hello world!"
+
+  def test_maybe_save_output_to_state_agent_transfer_scenario(self, caplog):
+    """Test realistic agent transfer scenario."""
+    # Scenario: Agent A transfers to Agent B, Agent B produces output
+    # Agent A should not save Agent B's output
+
+    # Set the LlmAgent logger to DEBUG level
+    llm_agent_logger = logging.getLogger(
+        "google_adk.google.adk.agents.llm_agent"
+    )
+    original_level = llm_agent_logger.level
+    llm_agent_logger.setLevel(logging.DEBUG)
+
+    try:
+      agent_a = LlmAgent(name="support_agent", output_key="support_result")
+      agent_b_event = create_test_event(
+          author="billing_agent", content_text="Your bill is $100"
+      )
+
+      with caplog.at_level("DEBUG"):
+        agent_a._LlmAgent__maybe_save_output_to_state(agent_b_event)
+
+      # Agent A should not save Agent B's output
+      assert len(agent_b_event.actions.state_delta) == 0
+      assert (
+          "Skipping output save for agent support_agent: event authored by"
+          " billing_agent"
+          in caplog.text
+      )
+    finally:
+      # Restore original logger level
+      llm_agent_logger.setLevel(original_level)
+
+  def test_maybe_save_output_to_state_case_sensitive_names(self, caplog):
+    """Test that agent name comparison is case-sensitive."""
+    # Set the LlmAgent logger to DEBUG level
+    llm_agent_logger = logging.getLogger(
+        "google_adk.google.adk.agents.llm_agent"
+    )
+    original_level = llm_agent_logger.level
+    llm_agent_logger.setLevel(logging.DEBUG)
+
+    try:
+      agent = LlmAgent(name="TestAgent", output_key="result")
+      event = create_test_event(
+          author="testagent", content_text="Test response"
+      )
+
+      with caplog.at_level("DEBUG"):
+        agent._LlmAgent__maybe_save_output_to_state(event)
+
+      # Should not save due to case mismatch
+      assert len(event.actions.state_delta) == 0
+      assert (
+          "Skipping output save for agent TestAgent: event authored by"
+          " testagent"
+          in caplog.text
+      )
+    finally:
+      # Restore original logger level
+      llm_agent_logger.setLevel(original_level)
+
+  @patch("google.adk.agents.llm_agent.logger")
+  def test_maybe_save_output_to_state_logging(self, mock_logger):
+    """Test that debug logging works correctly."""
+    agent = LlmAgent(name="agent1", output_key="result")
+    event = create_test_event(author="agent2", content_text="Test response")
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Should call logger.debug with correct parameters
+    mock_logger.debug.assert_called_once_with(
+        "Skipping output save for agent %s: event authored by %s",
+        "agent1",
+        "agent2",
+    )
+
+  @pytest.mark.parametrize("empty_content", ["", "  ", "\n"])
+  def test_maybe_save_output_to_state_handles_empty_final_chunk_with_schema(
+      self, empty_content
+  ):
+    """Tests that the agent correctly handles an empty final streaming chunk
+
+    when an output_schema is specified, preventing a crash.
+    """
+    # ARRANGE: Create an agent that expects a JSON output matching a schema.
+    agent = LlmAgent(
+        name="test_agent", output_key="result", output_schema=MockOutputSchema
+    )
+
+    # ARRANGE: Create a final event with empty or whitespace-only content.
+    # This simulates the final, empty chunk from a model's streaming response.
+    event = create_test_event(
+        author="test_agent", content_text=empty_content, is_final=True
+    )
+
+    # ACT: Call the method. The test's primary goal is to ensure this
+    # does NOT raise a pydantic.ValidationError, which it would have before the fix.
+    try:
+      agent._LlmAgent__maybe_save_output_to_state(event)
+    except Exception as e:
+      pytest.fail(f"The method unexpectedly raised an exception: {e}")
+
+    # ASSERT: Because the method should return early, the state_delta
+    # should remain empty.
+    assert len(event.actions.state_delta) == 0
+
+  @pytest.mark.asyncio
+  async def test_output_key_saved_when_before_agent_callback_short_circuits(
+      self,
+  ):
+    """Test that output_key is written to session state when
+    before_agent_callback short-circuits the agent."""
+
+    def cache_callback(callback_context: CallbackContext) -> types.Content:
+      return types.Content(parts=[types.Part.from_text(text="cached answer")])
+
+    agent = LlmAgent(
+        name="test_agent",
+        output_key="result",
+        before_agent_callback=cache_callback,
+    )
+
+    runner = testing_utils.InMemoryRunner(agent)
+    await runner.run_async("hello")
+
+    assert runner.session.state.get("result") == "cached answer"
+
+  def test_maybe_save_output_to_state_skips_function_response_only_event(self):
+    """Test that state_delta set by callback is not overwritten when event
+
+    only has function_response parts and no text.
+    """
+    agent = LlmAgent(name="test_agent", output_key="result")
+
+    # Simulate a function_response-only event (no text parts)
+    parts = [
+        types.Part(
+            function_response=types.FunctionResponse(
+                name="my_tool",
+                response={"status": "success", "data": [1, 2, 3]},
+            )
+        )
+    ]
+    content = types.Content(role="user", parts=parts)
+
+    event = Event(
+        invocation_id="test_invocation",
+        author="test_agent",
+        content=content,
+        actions=EventActions(
+            skip_summarization=True,
+            state_delta={"result": [1, 2, 3]},
+        ),
+    )
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # The callback-set value should be preserved, not overwritten with ""
+    assert event.actions.state_delta["result"] == [1, 2, 3]
+
+  def test_maybe_save_output_to_state_saves_empty_string_when_text_is_empty(
+      self,
+  ):
+    """Test that output is saved as empty string when part.text is explicitly empty."""
+    agent = LlmAgent(name="test_agent", output_key="result")
+
+    # Explicitly construct a part with empty string text
+    parts = [types.Part(text="")]
+    content = types.Content(role="model", parts=parts)
+    event = Event(
+        invocation_id="test_invocation",
+        author="test_agent",
+        content=content,
+        actions=EventActions(),
+    )
+
+    agent._LlmAgent__maybe_save_output_to_state(event)
+
+    # Assert key exists and value is empty string
+    assert "result" in event.actions.state_delta
+    assert not event.actions.state_delta["result"]

tests/unittests/agents/test_llm_agent_single_turn_subagents.py

@@ -0,0 +1,49 @@
+# 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.
+
+from google.adk.agents import LlmAgent
+
+
+def test_single_turn_managed_agent_is_wrapped_as_tool():
+  from google.adk.agents import ManagedAgent
+  from google.adk.tools.agent_tool import _SingleTurnAgentTool
+
+  managed = ManagedAgent(name='m', agent_id='a', mode='single_turn')
+  coordinator = LlmAgent(name='c', sub_agents=[managed])
+
+  assert any(
+      isinstance(t, _SingleTurnAgentTool) and t.agent is managed
+      for t in coordinator.tools
+  )
+
+
+def test_managed_agent_without_mode_is_not_wrapped():
+  from google.adk.agents import ManagedAgent
+  from google.adk.tools.agent_tool import _SingleTurnAgentTool
+
+  managed = ManagedAgent(name='m', agent_id='a')  # mode defaults to None
+  coordinator = LlmAgent(name='c', sub_agents=[managed])
+
+  assert not any(isinstance(t, _SingleTurnAgentTool) for t in coordinator.tools)
+
+
+def test_single_turn_managed_agent_excluded_from_transfer_targets():
+  from google.adk.agents import ManagedAgent
+  from google.adk.flows.llm_flows.agent_transfer import _get_transfer_targets
+
+  managed = ManagedAgent(name='m', agent_id='a', mode='single_turn')
+  coordinator = LlmAgent(name='c', sub_agents=[managed])
+
+  targets = _get_transfer_targets(coordinator)
+  assert managed not in targets

tests/unittests/agents/test_llm_agent_streaming_output.py

@@ -0,0 +1,120 @@
+# 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.
+
+"""Integration tests for LlmAgent output_key under streaming with tool calls."""
+
+from __future__ import annotations
+
+from typing import AsyncGenerator
+from unittest import mock
+
+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.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 import types
+import pytest
+
+
+def _text(text: str) -> types.Part:
+  return types.Part.from_text(text=text)
+
+
+def _call(name: str) -> types.Part:
+  return types.Part(function_call=types.FunctionCall(name=name, args={}))
+
+
+def _response(name: str) -> types.Part:
+  return types.Part(
+      function_response=types.FunctionResponse(name=name, response={"ok": True})
+  )
+
+
+def _event(
+    parts: list[types.Part], *, role: str = "model", partial: bool = False
+) -> Event:
+  return Event(
+      invocation_id="inv",
+      author="agent",
+      content=types.Content(role=role, parts=parts),
+      actions=EventActions(),
+      partial=partial,
+  )
+
+
+@pytest.mark.asyncio
+async def test_run_async_accumulates_text_around_tool_calls():
+  """Regression test for issue #5590.
+
+  Under StreamingMode.SSE with tools, an LlmAgent emits text in several
+  non-partial events: some carry text only, others carry text alongside a
+  function_call. Event.is_final_response() returns False for any event
+  with a function_call or function_response part, so the text on those
+  events was historically dropped from output_key — only the final
+  tool-free event's text was saved. Reporters measured ~60-70% loss.
+
+  Drive _run_async_impl with a stubbed _llm_flow that yields the canned
+  event sequence the streaming flow produces, merge each event's
+  state_delta into session state via session_service.append_event, and
+  assert that the user-visible session.state[output_key] contains every
+  non-partial text segment the agent emitted, in order.
+  """
+  canned = [
+      _event([_text("Intro one. ")], partial=True),
+      _event([_text("Intro two.")], partial=True),
+      _event([_text("Intro one. Intro two."), _call("t")]),
+      _event([_response("t")], role="user"),
+      _event([_text("Progress.")], partial=True),
+      _event([_text("Progress."), _call("t")]),
+      _event([_response("t")], role="user"),
+      _event([_text("Conclusion one. ")], partial=True),
+      _event([_text("Conclusion two.")], partial=True),
+      _event([_text("Conclusion one. Conclusion two.")]),
+  ]
+
+  class _FakeFlow:
+
+    async def run_async(
+        self, _ctx: InvocationContext
+    ) -> AsyncGenerator[Event, None]:
+      for event in canned:
+        yield event
+
+  agent = LlmAgent(name="agent", output_key="final_output")
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name="t", user_id="u", session_id="s"
+  )
+  ctx = InvocationContext(
+      invocation_id="inv",
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      run_config=RunConfig(),
+  )
+
+  with mock.patch.object(
+      type(agent),
+      "_llm_flow",
+      new_callable=mock.PropertyMock,
+      return_value=_FakeFlow(),
+  ):
+    async for event in agent._run_async_impl(ctx):
+      await session_service.append_event(session, event)
+
+  assert session.state["final_output"] == (
+      "Intro one. Intro two.Progress.Conclusion one. Conclusion two."
+  )

tests/unittests/agents/test_loop_agent.py

@@ -0,0 +1,292 @@
+# 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.
+
+"""Testings for the SequentialAgent."""
+
+from typing import AsyncGenerator
+from unittest.mock import patch
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.loop_agent import LoopAgent
+from google.adk.agents.loop_agent import LoopAgentState
+from google.adk.apps 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 import types
+import pytest
+from typing_extensions import override
+
+from .. import testing_utils
+
+END_OF_AGENT = testing_utils.END_OF_AGENT
+
+
+class _TestingAgent(BaseAgent):
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text=f'Hello, async {self.name}!')]
+        ),
+    )
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text=f'Hello, live {self.name}!')]
+        ),
+    )
+
+
+class _TestingAgentWithEscalateAction(BaseAgent):
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text=f'Hello, async {self.name}!')]
+        ),
+        actions=EventActions(escalate=True),
+    )
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text='I have done my job after escalation!!')]
+        ),
+    )
+
+
+async def _create_parent_invocation_context(
+    test_name: str, agent: BaseAgent, resumable: bool = False
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id=f'{test_name}_invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=resumable),
+  )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('resumable', [True, False])
+async def test_run_async(request: pytest.FixtureRequest, resumable: bool):
+  agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
+  loop_agent = LoopAgent(
+      name=f'{request.function.__name__}_test_loop_agent',
+      max_iterations=2,
+      sub_agents=[
+          agent,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, loop_agent, resumable=resumable
+  )
+  events = [e async for e in loop_agent.run_async(parent_ctx)]
+
+  simplified_events = testing_utils.simplify_resumable_app_events(events)
+  if resumable:
+    expected_events = [
+        (
+            loop_agent.name,
+            {'current_sub_agent': agent.name, 'times_looped': 0},
+        ),
+        (agent.name, f'Hello, async {agent.name}!'),
+        (
+            loop_agent.name,
+            {'current_sub_agent': agent.name, 'times_looped': 1},
+        ),
+        (agent.name, f'Hello, async {agent.name}!'),
+        (loop_agent.name, END_OF_AGENT),
+    ]
+  else:
+    expected_events = [
+        (agent.name, f'Hello, async {agent.name}!'),
+        (agent.name, f'Hello, async {agent.name}!'),
+    ]
+  assert simplified_events == expected_events
+
+
+@pytest.mark.asyncio
+async def test_resume_async(request: pytest.FixtureRequest):
+  agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
+  agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  loop_agent = LoopAgent(
+      name=f'{request.function.__name__}_test_loop_agent',
+      max_iterations=2,
+      sub_agents=[
+          agent_1,
+          agent_2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, loop_agent, resumable=True
+  )
+  parent_ctx.agent_states[loop_agent.name] = LoopAgentState(
+      current_sub_agent=agent_2.name, times_looped=1
+  ).model_dump(mode='json')
+
+  events = [e async for e in loop_agent.run_async(parent_ctx)]
+
+  simplified_events = testing_utils.simplify_resumable_app_events(events)
+  expected_events = [
+      (agent_2.name, f'Hello, async {agent_2.name}!'),
+      (loop_agent.name, END_OF_AGENT),
+  ]
+  assert simplified_events == expected_events
+
+
+@pytest.mark.asyncio
+async def test_run_async_skip_if_no_sub_agent(request: pytest.FixtureRequest):
+  loop_agent = LoopAgent(
+      name=f'{request.function.__name__}_test_loop_agent',
+      max_iterations=2,
+      sub_agents=[],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, loop_agent
+  )
+  events = [e async for e in loop_agent.run_async(parent_ctx)]
+  assert not events
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('resumable', [True, False])
+async def test_run_async_with_escalate_action(
+    request: pytest.FixtureRequest, resumable: bool
+):
+  non_escalating_agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_non_escalating_agent'
+  )
+  escalating_agent = _TestingAgentWithEscalateAction(
+      name=f'{request.function.__name__}_test_escalating_agent'
+  )
+  ignored_agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_ignored_agent'
+  )
+  loop_agent = LoopAgent(
+      name=f'{request.function.__name__}_test_loop_agent',
+      sub_agents=[non_escalating_agent, escalating_agent, ignored_agent],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, loop_agent, resumable=resumable
+  )
+  events = [e async for e in loop_agent.run_async(parent_ctx)]
+
+  simplified_events = testing_utils.simplify_resumable_app_events(events)
+
+  if resumable:
+    expected_events = [
+        (
+            loop_agent.name,
+            {
+                'current_sub_agent': non_escalating_agent.name,
+                'times_looped': 0,
+            },
+        ),
+        (
+            non_escalating_agent.name,
+            f'Hello, async {non_escalating_agent.name}!',
+        ),
+        (
+            loop_agent.name,
+            {'current_sub_agent': escalating_agent.name, 'times_looped': 0},
+        ),
+        (
+            escalating_agent.name,
+            f'Hello, async {escalating_agent.name}!',
+        ),
+        (
+            escalating_agent.name,
+            'I have done my job after escalation!!',
+        ),
+        (loop_agent.name, END_OF_AGENT),
+    ]
+  else:
+    expected_events = [
+        (
+            non_escalating_agent.name,
+            f'Hello, async {non_escalating_agent.name}!',
+        ),
+        (
+            escalating_agent.name,
+            f'Hello, async {escalating_agent.name}!',
+        ),
+        (
+            escalating_agent.name,
+            'I have done my job after escalation!!',
+        ),
+    ]
+  assert simplified_events == expected_events
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_pause_preserves_sub_agent_state(
+    request: pytest.FixtureRequest,
+):
+  """Test that the sub-agent state is preserved when the loop agent pauses."""
+  agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
+  loop_agent = LoopAgent(
+      name=f'{request.function.__name__}_test_loop_agent',
+      max_iterations=2,
+      sub_agents=[agent],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, loop_agent, resumable=True
+  )
+
+  # Set some dummy state for the sub-agent
+  parent_ctx.agent_states[agent.name] = {'some_key': 'some_value'}
+
+  # Mock should_pause_invocation to return True for the agent's event
+  def mock_should_pause(event):
+    return event.author == agent.name
+
+  with patch.object(
+      InvocationContext,
+      'should_pause_invocation',
+      side_effect=mock_should_pause,
+  ):
+    async for _ in loop_agent.run_async(parent_ctx):
+      pass  # Consume the async generator
+
+  # Verify that the sub-agent state was NOT reset
+  assert agent.name in parent_ctx.agent_states
+  assert parent_ctx.agent_states[agent.name] == {'some_key': 'some_value'}
+
+
+def test_deprecation_mentions_sub_agent_limitation():
+  with pytest.warns(DeprecationWarning, match='sub-agent'):
+    LoopAgent(name='deprecated_loop', sub_agents=[])

tests/unittests/agents/test_mcp_instruction_provider.py

@@ -0,0 +1,191 @@
+# 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.
+
+"""Unit tests for McpInstructionProvider."""
+
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.agents.mcp_instruction_provider import McpInstructionProvider
+from google.adk.agents.readonly_context import ReadonlyContext
+import pytest
+
+
+class TestMcpInstructionProvider:
+  """Unit tests for McpInstructionProvider."""
+
+  def setup_method(self):
+    """Sets up the test environment."""
+    self.connection_params = {"host": "localhost", "port": 8000}
+    self.prompt_name = "test_prompt"
+    self.mock_mcp_session_manager_cls = patch(
+        "google.adk.agents.mcp_instruction_provider.MCPSessionManager"
+    ).start()
+    self.mock_mcp_session_manager = (
+        self.mock_mcp_session_manager_cls.return_value
+    )
+    self.mock_session = MagicMock()
+    self.mock_session.list_prompts = AsyncMock()
+    self.mock_session.get_prompt = AsyncMock()
+    self.mock_mcp_session_manager.create_session = AsyncMock(
+        return_value=self.mock_session
+    )
+    self.provider = McpInstructionProvider(
+        self.connection_params, self.prompt_name
+    )
+
+  @pytest.mark.asyncio
+  async def test_call_success_no_args(self):
+    """Tests __call__ with a prompt that has no arguments."""
+    mock_prompt = MagicMock()
+    mock_prompt.name = self.prompt_name
+    mock_prompt.arguments = None
+    self.mock_session.list_prompts.return_value = MagicMock(
+        prompts=[mock_prompt]
+    )
+
+    mock_msg1 = MagicMock()
+    mock_msg1.content.type = "text"
+    mock_msg1.content.text = "instruction part 1. "
+    mock_msg2 = MagicMock()
+    mock_msg2.content.type = "text"
+    mock_msg2.content.text = "instruction part 2"
+    self.mock_session.get_prompt.return_value = MagicMock(
+        messages=[mock_msg1, mock_msg2]
+    )
+
+    mock_invocation_context = MagicMock()
+    mock_invocation_context.session.state = {}
+    context = ReadonlyContext(mock_invocation_context)
+
+    # Call
+    instruction = await self.provider(context)
+
+    # Assert
+    assert instruction == "instruction part 1. instruction part 2"
+    self.mock_session.get_prompt.assert_called_once_with(
+        self.prompt_name, arguments={}
+    )
+
+  @pytest.mark.asyncio
+  async def test_call_success_with_args(self):
+    """Tests __call__ with a prompt that has arguments."""
+    mock_arg1 = MagicMock()
+    mock_arg1.name = "arg1"
+    mock_prompt = MagicMock()
+    mock_prompt.name = self.prompt_name
+    mock_prompt.arguments = [mock_arg1]
+    self.mock_session.list_prompts.return_value = MagicMock(
+        prompts=[mock_prompt]
+    )
+
+    mock_msg = MagicMock()
+    mock_msg.content.type = "text"
+    mock_msg.content.text = "instruction with arg1"
+    self.mock_session.get_prompt.return_value = MagicMock(messages=[mock_msg])
+
+    mock_invocation_context = MagicMock()
+    mock_invocation_context.session.state = {"arg1": "value1", "arg2": "value2"}
+    context = ReadonlyContext(mock_invocation_context)
+
+    instruction = await self.provider(context)
+
+    assert instruction == "instruction with arg1"
+    self.mock_session.get_prompt.assert_called_once_with(
+        self.prompt_name, arguments={"arg1": "value1"}
+    )
+
+  @pytest.mark.asyncio
+  async def test_call_prompt_not_found_in_list_prompts(self):
+    """Tests __call__ when list_prompts doesn't return the prompt."""
+    self.mock_session.list_prompts.return_value = MagicMock(prompts=[])
+
+    mock_msg = MagicMock()
+    mock_msg.content.type = "text"
+    mock_msg.content.text = "instruction"
+    self.mock_session.get_prompt.return_value = MagicMock(messages=[mock_msg])
+
+    mock_invocation_context = MagicMock()
+    mock_invocation_context.session.state = {"arg1": "value1"}
+    context = ReadonlyContext(mock_invocation_context)
+
+    instruction = await self.provider(context)
+
+    assert instruction == "instruction"
+    self.mock_session.get_prompt.assert_called_once_with(
+        self.prompt_name, arguments={}
+    )
+
+  @pytest.mark.asyncio
+  async def test_call_get_prompt_returns_no_messages(self):
+    """Tests __call__ when get_prompt returns no messages."""
+    # Setup mocks
+    self.mock_session.list_prompts.return_value = MagicMock(prompts=[])
+    self.mock_session.get_prompt.return_value = MagicMock(messages=[])
+
+    mock_invocation_context = MagicMock()
+    mock_invocation_context.session.state = {}
+    context = ReadonlyContext(mock_invocation_context)
+
+    # Call and assert
+    with pytest.raises(
+        ValueError, match="Failed to load MCP prompt 'test_prompt'."
+    ):
+      await self.provider(context)
+
+    # Assert
+    self.mock_session.get_prompt.assert_called_once_with(
+        self.prompt_name, arguments={}
+    )
+
+  @pytest.mark.asyncio
+  async def test_call_ignore_non_text_messages(self):
+    """Tests __call__ ignores non-text messages."""
+    # Setup mocks
+    mock_prompt = MagicMock()
+    mock_prompt.name = self.prompt_name
+    mock_prompt.arguments = None
+    self.mock_session.list_prompts.return_value = MagicMock(
+        prompts=[mock_prompt]
+    )
+
+    mock_msg1 = MagicMock()
+    mock_msg1.content.type = "text"
+    mock_msg1.content.text = "instruction part 1. "
+
+    mock_msg2 = MagicMock()
+    mock_msg2.content.type = "image"
+    mock_msg2.content.text = "ignored"
+
+    mock_msg3 = MagicMock()
+    mock_msg3.content.type = "text"
+    mock_msg3.content.text = "instruction part 2"
+
+    self.mock_session.get_prompt.return_value = MagicMock(
+        messages=[mock_msg1, mock_msg2, mock_msg3]
+    )
+
+    mock_invocation_context = MagicMock()
+    mock_invocation_context.session.state = {}
+    context = ReadonlyContext(mock_invocation_context)
+
+    # Call
+    instruction = await self.provider(context)
+
+    # Assert
+    assert instruction == "instruction part 1. instruction part 2"
+    self.mock_session.get_prompt.assert_called_once_with(
+        self.prompt_name, arguments={}
+    )

tests/unittests/agents/test_model_callback_chain.py

@@ -0,0 +1,250 @@
+# 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.
+
+from enum import Enum
+from functools import partial
+from typing import Any
+from typing import List
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from .. import testing_utils
+
+
+class CallbackType(Enum):
+  SYNC = 1
+  ASYNC = 2
+
+
+async def mock_async_before_cb_side_effect(
+    callback_context: CallbackContext,
+    llm_request: LlmRequest,
+    ret_value=None,
+):
+  if ret_value:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=ret_value)]
+        )
+    )
+  return None
+
+
+def mock_sync_before_cb_side_effect(
+    callback_context: CallbackContext,
+    llm_request: LlmRequest,
+    ret_value=None,
+):
+  if ret_value:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=ret_value)]
+        )
+    )
+  return None
+
+
+async def mock_async_after_cb_side_effect(
+    callback_context: CallbackContext,
+    llm_response: LlmResponse,
+    ret_value=None,
+):
+  if ret_value:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=ret_value)]
+        )
+    )
+  return None
+
+
+def mock_sync_after_cb_side_effect(
+    callback_context: CallbackContext,
+    llm_response: LlmResponse,
+    ret_value=None,
+):
+  if ret_value:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=ret_value)]
+        )
+    )
+  return None
+
+
+CALLBACK_PARAMS = [
+    pytest.param(
+        [
+            (None, CallbackType.SYNC),
+            ("callback_2_response", CallbackType.ASYNC),
+            ("callback_3_response", CallbackType.SYNC),
+            (None, CallbackType.ASYNC),
+        ],
+        "callback_2_response",
+        [1, 1, 0, 0],
+        id="middle_async_callback_returns",
+    ),
+    pytest.param(
+        [
+            (None, CallbackType.SYNC),
+            (None, CallbackType.ASYNC),
+            (None, CallbackType.SYNC),
+            (None, CallbackType.ASYNC),
+        ],
+        "model_response",
+        [1, 1, 1, 1],
+        id="all_callbacks_return_none",
+    ),
+    pytest.param(
+        [
+            ("callback_1_response", CallbackType.SYNC),
+            ("callback_2_response", CallbackType.ASYNC),
+        ],
+        "callback_1_response",
+        [1, 0],
+        id="first_sync_callback_returns",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "callbacks, expected_response, expected_calls",
+    CALLBACK_PARAMS,
+)
+@pytest.mark.asyncio
+async def test_before_model_callbacks_chain(
+    callbacks: List[tuple[str, int]],
+    expected_response: str,
+    expected_calls: List[int],
+):
+  responses = ["model_response"]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+
+  mock_cbs = []
+  for response, callback_type in callbacks:
+
+    if callback_type == CallbackType.ASYNC:
+      mock_cb = mock.AsyncMock(
+          side_effect=partial(
+              mock_async_before_cb_side_effect, ret_value=response
+          )
+      )
+    else:
+      mock_cb = mock.Mock(
+          side_effect=partial(
+              mock_sync_before_cb_side_effect, ret_value=response
+          )
+      )
+    mock_cbs.append(mock_cb)
+  # Create agent with multiple callbacks
+  agent = Agent(
+      name="root_agent",
+      model=mock_model,
+      before_model_callback=[mock_cb for mock_cb in mock_cbs],
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  result = await runner.run_async_with_new_session("test")
+  assert testing_utils.simplify_events(result) == [
+      ("root_agent", expected_response),
+  ]
+
+  # Assert that the callbacks were called the expected number of times
+  for i, mock_cb in enumerate(mock_cbs):
+    expected_calls_count = expected_calls[i]
+    if expected_calls_count == 1:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited_once()
+      else:
+        mock_cb.assert_called_once()
+    elif expected_calls_count == 0:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_not_awaited()
+      else:
+        mock_cb.assert_not_called()
+    else:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited(expected_calls_count)
+      else:
+        mock_cb.assert_called(expected_calls_count)
+
+
+@pytest.mark.parametrize(
+    "callbacks, expected_response, expected_calls",
+    CALLBACK_PARAMS,
+)
+@pytest.mark.asyncio
+async def test_after_model_callbacks_chain(
+    callbacks: List[tuple[str, int]],
+    expected_response: str,
+    expected_calls: List[int],
+):
+  responses = ["model_response"]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+
+  mock_cbs = []
+  for response, callback_type in callbacks:
+
+    if callback_type == CallbackType.ASYNC:
+      mock_cb = mock.AsyncMock(
+          side_effect=partial(
+              mock_async_after_cb_side_effect, ret_value=response
+          )
+      )
+    else:
+      mock_cb = mock.Mock(
+          side_effect=partial(
+              mock_sync_after_cb_side_effect, ret_value=response
+          )
+      )
+    mock_cbs.append(mock_cb)
+  # Create agent with multiple callbacks
+  agent = Agent(
+      name="root_agent",
+      model=mock_model,
+      after_model_callback=[mock_cb for mock_cb in mock_cbs],
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  result = await runner.run_async_with_new_session("test")
+  assert testing_utils.simplify_events(result) == [
+      ("root_agent", expected_response),
+  ]
+
+  # Assert that the callbacks were called the expected number of times
+  for i, mock_cb in enumerate(mock_cbs):
+    expected_calls_count = expected_calls[i]
+    if expected_calls_count == 1:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited_once()
+      else:
+        mock_cb.assert_called_once()
+    elif expected_calls_count == 0:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_not_awaited()
+      else:
+        mock_cb.assert_not_called()
+    else:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited(expected_calls_count)
+      else:
+        mock_cb.assert_called(expected_calls_count)

tests/unittests/agents/test_output_key_visibility.py

@@ -0,0 +1,180 @@
+# 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 LlmAgent output_key visibility in callbacks."""
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.auto_flow import AutoFlow
+from google.genai import types
+import pytest
+from pytest_mock import MockerFixture
+
+from .. import testing_utils
+
+# Standard MockModel will be used instead of SafeMockModel
+
+
+@pytest.mark.asyncio
+async def test_output_key_visibility_in_after_agent_callback():
+  """Test that output_key state delta is visible in after_agent_callback."""
+  mock_response = "Hello! How can I help you?"
+  mock_model = testing_utils.MockModel.create(responses=[mock_response])
+
+  callback_called = False
+  captured_state_value = None
+  captured_session_state_value = None
+
+  async def check_output_key(callback_context: CallbackContext):
+    nonlocal callback_called, captured_state_value, captured_session_state_value
+    callback_called = True
+    captured_state_value = callback_context.state.get("result", "NOT_FOUND")
+    captured_session_state_value = callback_context.session.state.get(
+        "result", "NOT_IN_RAW"
+    )
+
+  agent = LlmAgent(
+      name="my_agent",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      output_key="result",
+      after_agent_callback=check_output_key,
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent=agent)
+
+  events = await runner.run_async(new_message="hello")
+
+  assert callback_called, "Callback was not called"
+
+  assert (
+      captured_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_state_value}"
+  assert (
+      captured_session_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_session_state_value}"
+
+
+@pytest.mark.asyncio
+async def test_output_key_visibility_in_run_live(mocker: MockerFixture):
+  """Test that output_key state delta is visible in after_agent_callback in run_live."""
+  mock_response = "Hello! How can I help you?"
+  mock_model = testing_utils.MockModel.create(responses=[mock_response])
+
+  callback_called = False
+  captured_state_value = None
+  captured_session_state_value = None
+
+  async def check_output_key(callback_context: CallbackContext):
+    nonlocal callback_called, captured_state_value, captured_session_state_value
+    callback_called = True
+    captured_state_value = callback_context.state.get("result", "NOT_FOUND")
+    captured_session_state_value = callback_context.session.state.get(
+        "result", "NOT_IN_RAW"
+    )
+
+  agent = LlmAgent(
+      name="my_agent",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      output_key="result",
+      after_agent_callback=check_output_key,
+  )
+
+  async def mock_auto_flow_run_live(self, ctx):
+    yield Event(
+        id=Event.new_id(),
+        invocation_id=ctx.invocation_id,
+        author=ctx.agent.name,
+        content=types.Content(parts=[types.Part(text=mock_response)]),
+    )
+
+  mocker.patch.object(AutoFlow, "run_live", mock_auto_flow_run_live)
+
+  runner = testing_utils.InMemoryRunner(root_agent=agent)
+  live_queue = LiveRequestQueue()
+
+  agen = runner.runner.run_live(
+      user_id="test_user",
+      session_id=runner.session.id,
+      live_request_queue=live_queue,
+  )
+
+  # Send a message to trigger the agent
+  live_queue.send_content(
+      types.Content(role="user", parts=[types.Part(text="hello")])
+  )
+
+  live_queue.close()
+
+  async for event in agen:
+    pass
+
+  assert callback_called, "Callback was not called"
+  assert (
+      captured_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_state_value}"
+  assert (
+      captured_session_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_session_state_value}"
+
+
+@pytest.mark.asyncio
+async def test_output_key_visibility_in_sequential_agent():
+  """Test that output_key state delta is visible in next agent's before_agent_callback."""
+  mock_response = "Hello from agent 1!"
+  mock_model = testing_utils.MockModel.create(
+      responses=[mock_response, "Hello from agent 2!"]
+  )
+
+  callback_called = False
+  captured_session_state_value = None
+
+  async def check_before_agent(callback_context: CallbackContext):
+    nonlocal callback_called, captured_session_state_value
+    callback_called = True
+    captured_session_state_value = callback_context.session.state.get(
+        "result", "NOT_FOUND"
+    )
+
+  agent_1 = LlmAgent(
+      name="agent_1",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      output_key="result",
+  )
+
+  agent_2 = LlmAgent(
+      name="agent_2",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      before_agent_callback=check_before_agent,
+  )
+
+  sequential_agent = SequentialAgent(
+      name="seq_agent",
+      sub_agents=[agent_1, agent_2],
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent=sequential_agent)
+
+  events = await runner.run_async(new_message="hello")
+
+  assert callback_called, "Callback was not called"
+  assert (
+      captured_session_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_session_state_value}"

tests/unittests/agents/test_parallel_agent.py

@@ -0,0 +1,416 @@
+# 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 ParallelAgent."""
+
+import asyncio
+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 _merge_agent_run_pre_3_11
+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
+import pytest
+from typing_extensions import override
+
+
+class _TestingAgent(BaseAgent):
+
+  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=types.Content(
+            parts=[types.Part(text=f'Hello, async {self.name}!')]
+        ),
+    )
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    await asyncio.sleep(self.delay)
+    yield self.event(ctx)
+    if ctx.is_resumable:
+      ctx.set_agent_state(self.name, end_of_agent=True)
+
+
+async def _create_parent_invocation_context(
+    test_name: str, agent: BaseAgent, is_resumable: bool = False
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id=f'{test_name}_invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=is_resumable),
+  )
+
+
+@pytest.mark.asyncio
+@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,
+  )
+  agent2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[
+          agent1,
+          agent2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent, is_resumable=is_resumable
+  )
+  events = [e async for e in parallel_agent.run_async(parent_ctx)]
+
+  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
+@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,
+  )
+  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,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent, is_resumable=is_resumable
+  )
+  events = [e async for e in parallel_agent.run_async(parent_ctx)]
+
+  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,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent, is_resumable=True
+  )
+  parent_ctx.agent_states[parallel_agent.name] = BaseAgentState().model_dump(
+      mode='json'
+  )
+  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):
+  """Mock agent for testing."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    for _ in range(0, 3):
+      event = self.event(ctx)
+      yield event
+      # Check that the event was processed by the consumer.
+      assert event.custom_metadata is not None
+      assert event.custom_metadata['processed']
+
+
+@pytest.mark.asyncio
+async def test_generating_one_event_per_agent_at_once(
+    request: pytest.FixtureRequest,
+):
+  # This test is to verify that the parallel agent won't generate more than one
+  # event per agent at a time.
+  agent1 = _TestingAgentWithMultipleEvents(
+      name=f'{request.function.__name__}_test_agent_1'
+  )
+  agent2 = _TestingAgentWithMultipleEvents(
+      name=f'{request.function.__name__}_test_agent_2'
+  )
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[
+          agent1,
+          agent2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent
+  )
+
+  agen = parallel_agent.run_async(parent_ctx)
+  async for event in agen:
+    event.custom_metadata = {'processed': True}
+    # Asserts on event are done in _TestingAgentWithMultipleEvents.
+
+
+@pytest.mark.asyncio
+async def test_run_async_skip_if_no_sub_agent(request: pytest.FixtureRequest):
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent
+  )
+  events = [e async for e in parallel_agent.run_async(parent_ctx)]
+  assert not events
+
+
+class _TestingAgentWithException(_TestingAgent):
+  """Mock agent for testing."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield self.event(ctx)
+    raise Exception()
+
+
+class _TestingAgentInfiniteEvents(_TestingAgent):
+  """Mock agent for testing."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    while True:
+      yield self.event(ctx)
+
+
+@pytest.mark.asyncio
+async def test_stop_agent_if_sub_agent_fails(
+    request: pytest.FixtureRequest,
+):
+  # This test is to verify that the parallel agent and subagents will all stop
+  # processing and throw exception to top level runner in case of exception.
+  agent1 = _TestingAgentWithException(
+      name=f'{request.function.__name__}_test_agent_1'
+  )
+  agent2 = _TestingAgentInfiniteEvents(
+      name=f'{request.function.__name__}_test_agent_2'
+  )
+  parallel_agent = ParallelAgent(
+      name=f'{request.function.__name__}_test_parallel_agent',
+      sub_agents=[
+          agent1,
+          agent2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, parallel_agent
+  )
+
+  agen = parallel_agent.run_async(parent_ctx)
+  # We expect to receive an exception from one of subagents.
+  # The exception should be propagated to root agent and other subagents.
+  # Otherwise we'll have an infinite loop.
+  with pytest.raises(Exception):
+    async for _ in agen:
+      # The infinite agent could iterate a few times depending on scheduling.
+      pass
+
+
+async def _slow_agent_with_cleanup_delay():
+  """Async generator that sleeps in its finally block to simulate cleanup."""
+  try:
+    await asyncio.sleep(10)
+    yield 'slow-event'
+  finally:
+    await asyncio.sleep(0.05)
+
+
+async def _failing_agent():
+  """Async generator that raises after a short delay."""
+  await asyncio.sleep(0.01)
+  raise ValueError('simulated sub-agent failure')
+  yield  # pragma: no cover
+
+
+@pytest.mark.asyncio
+async def test_merge_agent_run_pre_3_11_no_aclose_error_on_failure():
+  """Regression test for Python 3.10 RuntimeError: aclose() already running.
+
+  _merge_agent_run_pre_3_11 must await all cancelled tasks before returning so
+  that generators are fully released before the caller invokes aclose() on them.
+  """
+  agent_runs = [_slow_agent_with_cleanup_delay(), _failing_agent()]
+
+  with pytest.raises(ValueError, match='simulated sub-agent failure'):
+    async for _ in _merge_agent_run_pre_3_11(agent_runs):
+      pass
+
+  # If tasks were not properly awaited, aclose() on a still-running generator
+  # would raise RuntimeError here.
+  for agen in agent_runs:
+    await agen.aclose()
+
+
+def test_deprecation_mentions_sub_agent_limitation():
+  with pytest.warns(DeprecationWarning, match='sub-agent'):
+    ParallelAgent(name='deprecated_parallel', sub_agents=[])

tests/unittests/agents/test_readonly_context.py

@@ -0,0 +1,59 @@
+# 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.
+
+from types import MappingProxyType
+from unittest.mock import MagicMock
+
+from google.adk.agents.readonly_context import ReadonlyContext
+import pytest
+
+
+@pytest.fixture
+def mock_invocation_context():
+  mock_context = MagicMock()
+  mock_context.invocation_id = "test-invocation-id"
+  mock_context.agent.name = "test-agent-name"
+  mock_context.session.state = {"key1": "value1", "key2": "value2"}
+  mock_context.user_id = "test-user-id"
+  return mock_context
+
+
+def test_invocation_id(mock_invocation_context):
+  readonly_context = ReadonlyContext(mock_invocation_context)
+  assert readonly_context.invocation_id == "test-invocation-id"
+
+
+def test_agent_name(mock_invocation_context):
+  readonly_context = ReadonlyContext(mock_invocation_context)
+  assert readonly_context.agent_name == "test-agent-name"
+
+
+def test_agent_name_without_agent(mock_invocation_context):
+  mock_invocation_context.agent = None
+  readonly_context = ReadonlyContext(mock_invocation_context)
+  assert readonly_context.agent_name == "unknown"
+
+
+def test_state_content(mock_invocation_context):
+  readonly_context = ReadonlyContext(mock_invocation_context)
+  state = readonly_context.state
+
+  assert isinstance(state, MappingProxyType)
+  assert state["key1"] == "value1"
+  assert state["key2"] == "value2"
+
+
+def test_user_id(mock_invocation_context):
+  readonly_context = ReadonlyContext(mock_invocation_context)
+  assert readonly_context.user_id == "test-user-id"

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),
+  ]

tests/unittests/agents/test_run_config.py

@@ -0,0 +1,139 @@
+# 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.
+
+import sys
+from unittest.mock import ANY
+from unittest.mock import patch
+import warnings
+
+from google.adk.agents.run_config import RunConfig
+from google.genai import types
+import pytest
+
+
+def test_validate_max_llm_calls_valid():
+  value = RunConfig.validate_max_llm_calls(100)
+  assert value == 100
+
+
+def test_validate_max_llm_calls_negative():
+  with patch("google.adk.agents.run_config.logger.warning") as mock_warning:
+    value = RunConfig.validate_max_llm_calls(-1)
+    mock_warning.assert_called_once_with(ANY)
+    assert value == -1
+
+
+def test_validate_max_llm_calls_warns_on_zero():
+  with patch("google.adk.agents.run_config.logger.warning") as mock_warning:
+    value = RunConfig.validate_max_llm_calls(0)
+    mock_warning.assert_called_once_with(ANY)
+    assert value == 0
+
+
+def test_validate_max_llm_calls_too_large():
+  with pytest.raises(
+      ValueError, match=f"max_llm_calls should be less than {sys.maxsize}."
+  ):
+    RunConfig.validate_max_llm_calls(sys.maxsize)
+
+
+def test_audio_transcription_configs_are_not_shared_between_instances():
+  config1 = RunConfig()
+  config2 = RunConfig()
+
+  # Validate output_audio_transcription
+  assert config1.output_audio_transcription is not None
+  assert config2.output_audio_transcription is not None
+  assert (
+      config1.output_audio_transcription
+      is not config2.output_audio_transcription
+  )
+
+  # Validate input_audio_transcription
+  assert config1.input_audio_transcription is not None
+  assert config2.input_audio_transcription is not None
+  assert (
+      config1.input_audio_transcription is not config2.input_audio_transcription
+  )
+
+
+def test_response_modalities_accepts_enum():
+  config = RunConfig(response_modalities=[types.Modality.AUDIO])
+  assert config.response_modalities == [types.Modality.AUDIO]
+  assert isinstance(config.response_modalities[0], types.Modality)
+
+
+def test_response_modalities_coerces_string_to_enum():
+  config = RunConfig(response_modalities=["AUDIO"])
+  assert config.response_modalities == [types.Modality.AUDIO]
+  assert isinstance(config.response_modalities[0], types.Modality)
+
+
+def test_response_modalities_coerces_lowercase_string_to_enum():
+  config = RunConfig(response_modalities=["audio"])
+  assert config.response_modalities == [types.Modality.AUDIO]
+  assert isinstance(config.response_modalities[0], types.Modality)
+
+
+def test_response_modalities_serialization_no_warning():
+  config = RunConfig(response_modalities=[types.Modality.AUDIO])
+  live_config = types.LiveConnectConfig()
+  live_config.response_modalities = config.response_modalities
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+    live_config.model_dump()
+    pydantic_warnings = [
+        x for x in w if "PydanticSerializationUnexpectedValue" in str(x.message)
+    ]
+    assert len(pydantic_warnings) == 0
+
+
+def test_avatar_config_initialization():
+  custom_avatar = types.CustomizedAvatar(
+      image_mime_type="image/jpeg", image_data=b"image_bytes"
+  )
+  avatar_config = types.AvatarConfig(
+      audio_bitrate_bps=128000,
+      video_bitrate_bps=1000000,
+      customized_avatar=custom_avatar,
+  )
+  run_config = RunConfig(avatar_config=avatar_config)
+
+  assert run_config.avatar_config == avatar_config
+  assert run_config.avatar_config.customized_avatar == custom_avatar
+  assert (
+      run_config.avatar_config.customized_avatar.image_mime_type == "image/jpeg"
+  )
+  assert run_config.avatar_config.customized_avatar.image_data == b"image_bytes"
+
+
+def test_avatar_config_with_name():
+  avatar_config = types.AvatarConfig(
+      audio_bitrate_bps=128000,
+      video_bitrate_bps=1000000,
+      avatar_name="test_avatar",
+  )
+  run_config = RunConfig(avatar_config=avatar_config)
+
+  assert run_config.avatar_config == avatar_config
+  assert run_config.avatar_config.avatar_name == "test_avatar"
+  assert run_config.avatar_config.customized_avatar is None
+
+
+def test_model_input_context_accepts_transient_contents():
+  context_content = types.UserContent("Relevant context for this turn")
+
+  run_config = RunConfig(model_input_context=[context_content])
+
+  assert run_config.model_input_context == [context_content]

tests/unittests/agents/test_sequential_agent.py

@@ -0,0 +1,208 @@
+# 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.
+
+"""Testings for the SequentialAgent."""
+
+from typing import AsyncGenerator
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.agents.sequential_agent import SequentialAgentState
+from google.adk.apps import ResumabilityConfig
+from google.adk.events.event import Event
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.genai import types
+import pytest
+from typing_extensions import override
+
+
+class _TestingAgent(BaseAgent):
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text=f'Hello, async {self.name}!')]
+        ),
+    )
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            parts=[types.Part(text=f'Hello, live {self.name}!')]
+        ),
+    )
+
+
+async def _create_parent_invocation_context(
+    test_name: str, agent: BaseAgent, resumable: bool = False
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id=f'{test_name}_invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      resumability_config=ResumabilityConfig(is_resumable=resumable),
+  )
+
+
+@pytest.mark.asyncio
+async def test_run_async(request: pytest.FixtureRequest):
+  agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
+  agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  sequential_agent = SequentialAgent(
+      name=f'{request.function.__name__}_test_agent',
+      sub_agents=[
+          agent_1,
+          agent_2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, sequential_agent
+  )
+  events = [e async for e in sequential_agent.run_async(parent_ctx)]
+
+  assert len(events) == 2
+  assert events[0].author == agent_1.name
+  assert events[1].author == agent_2.name
+  assert events[0].content.parts[0].text == f'Hello, async {agent_1.name}!'
+  assert events[1].content.parts[0].text == f'Hello, async {agent_2.name}!'
+
+
+@pytest.mark.asyncio
+async def test_run_async_skip_if_no_sub_agent(request: pytest.FixtureRequest):
+  sequential_agent = SequentialAgent(
+      name=f'{request.function.__name__}_test_agent',
+      sub_agents=[],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, sequential_agent
+  )
+  events = [e async for e in sequential_agent.run_async(parent_ctx)]
+
+  assert not events
+
+
+@pytest.mark.asyncio
+async def test_run_async_with_resumability(request: pytest.FixtureRequest):
+  agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
+  agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  sequential_agent = SequentialAgent(
+      name=f'{request.function.__name__}_test_agent',
+      sub_agents=[
+          agent_1,
+          agent_2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, sequential_agent, resumable=True
+  )
+  events = [e async for e in sequential_agent.run_async(parent_ctx)]
+
+  # 5 events:
+  # 1. SequentialAgent checkpoint event for agent 1
+  # 2. Agent 1 event
+  # 3. SequentialAgent checkpoint event for agent 2
+  # 4. Agent 2 event
+  # 5. SequentialAgent final checkpoint event
+  assert len(events) == 5
+  assert events[0].author == sequential_agent.name
+  assert not events[0].actions.end_of_agent
+  assert events[0].actions.agent_state['current_sub_agent'] == agent_1.name
+
+  assert events[1].author == agent_1.name
+  assert events[1].content.parts[0].text == f'Hello, async {agent_1.name}!'
+
+  assert events[2].author == sequential_agent.name
+  assert not events[2].actions.end_of_agent
+  assert events[2].actions.agent_state['current_sub_agent'] == agent_2.name
+
+  assert events[3].author == agent_2.name
+  assert events[3].content.parts[0].text == f'Hello, async {agent_2.name}!'
+
+  assert events[4].author == sequential_agent.name
+  assert events[4].actions.end_of_agent
+
+
+@pytest.mark.asyncio
+async def test_resume_async(request: pytest.FixtureRequest):
+  agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
+  agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  sequential_agent = SequentialAgent(
+      name=f'{request.function.__name__}_test_agent',
+      sub_agents=[
+          agent_1,
+          agent_2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, sequential_agent, resumable=True
+  )
+  parent_ctx.agent_states[sequential_agent.name] = SequentialAgentState(
+      current_sub_agent=agent_2.name
+  ).model_dump(mode='json')
+
+  events = [e async for e in sequential_agent.run_async(parent_ctx)]
+
+  # 2 events:
+  # 1. Agent 2 event
+  # 2. SequentialAgent final checkpoint event
+  assert len(events) == 2
+  assert events[0].author == agent_2.name
+  assert events[0].content.parts[0].text == f'Hello, async {agent_2.name}!'
+
+  assert events[1].author == sequential_agent.name
+  assert events[1].actions.end_of_agent
+
+
+@pytest.mark.asyncio
+async def test_run_live(request: pytest.FixtureRequest):
+  agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
+  agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
+  sequential_agent = SequentialAgent(
+      name=f'{request.function.__name__}_test_agent',
+      sub_agents=[
+          agent_1,
+          agent_2,
+      ],
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, sequential_agent
+  )
+  events = [e async for e in sequential_agent.run_live(parent_ctx)]
+
+  assert len(events) == 2
+  assert events[0].author == agent_1.name
+  assert events[1].author == agent_2.name
+  assert events[0].content.parts[0].text == f'Hello, live {agent_1.name}!'
+  assert events[1].content.parts[0].text == f'Hello, live {agent_2.name}!'
+
+
+def test_deprecation_mentions_sub_agent_limitation():
+  with pytest.warns(DeprecationWarning, match='sub-agent'):
+    SequentialAgent(name='deprecated_sequential', sub_agents=[])