Merge branch 'main' into expose_credential_service

tests/unittests/agents/test_agent_clone.py

@@ -0,0 +1,379 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Testings for the clone functionality of agents."""
+
+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 non-existent 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")
+  seq_cloned = seq_original.clone()
+  assert isinstance(seq_cloned, SequentialAgent)
+
+  # Test ParallelAgent
+  par_original = ParallelAgent(name="par_test")
+  par_cloned = par_original.clone()
+  assert isinstance(par_cloned, ParallelAgent)
+
+  # Test LoopAgent
+  loop_original = LoopAgent(name="loop_test")
+  loop_cloned = loop_original.clone()
+  assert isinstance(loop_cloned, LoopAgent)
+
+
+def test_clone_with_agent_specific_fields():
+  # Test LoopAgent
+  loop_original = LoopAgent(name="loop_test")
+  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"
+
+
+if __name__ == "__main__":
+  # Run a specific test for debugging
+  test_three_level_nested_agent()

tests/unittests/agents/test_base_agent.py

@@ -26,6 +26,8 @@ from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events import Event
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.plugins.plugin_manager import PluginManager
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.genai import types
 import pytest
@@ -83,6 +85,35 @@ async def _async_after_agent_callback_append_agent_reply(
   )
 
 
+class MockPlugin(BasePlugin):
+  before_agent_text = 'before_agent_text from MockPlugin'
+  after_agent_text = 'after_agent_text from MockPlugin'
+
+  def __init__(self, name='mock_plugin'):
+    self.name = name
+    self.enable_before_agent_callback = False
+    self.enable_after_agent_callback = False
+
+  async def before_agent_callback(
+      self, *, agent: BaseAgent, callback_context: CallbackContext
+  ) -> Optional[types.Content]:
+    if not self.enable_before_agent_callback:
+      return None
+    return types.Content(parts=[types.Part(text=self.before_agent_text)])
+
+  async def after_agent_callback(
+      self, *, agent: BaseAgent, callback_context: CallbackContext
+  ) -> Optional[types.Content]:
+    if not self.enable_after_agent_callback:
+      return None
+    return types.Content(parts=[types.Part(text=self.after_agent_text)])
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
 class _IncompleteAgent(BaseAgent):
   pass
 
@@ -113,7 +144,10 @@ class _TestingAgent(BaseAgent):
 
 
 async def _create_parent_invocation_context(
-    test_name: str, agent: BaseAgent, branch: Optional[str] = None
+    test_name: str,
+    agent: BaseAgent,
+    branch: Optional[str] = None,
+    plugins: list[BasePlugin] = [],
 ) -> InvocationContext:
   session_service = InMemorySessionService()
   session = await session_service.create_session(
@@ -125,6 +159,7 @@ async def _create_parent_invocation_context(
       agent=agent,
       session=session,
       session_service=session_service,
+      plugin_manager=PluginManager(plugins=plugins),
   )
 
 
@@ -190,6 +225,36 @@ async def test_run_async_before_agent_callback_noop(
   spy_run_async_impl.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_run_async_before_agent_callback_use_plugin(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+    mock_plugin: MockPlugin,
+):
+  """Test that the before agent callback uses the plugin response if both plugin callback and canonical agent callbacks are present."""
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      before_agent_callback=_before_agent_callback_bypass_agent,
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, agent, plugins=[mock_plugin]
+  )
+  mock_plugin.enable_before_agent_callback = True
+  spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
+  spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_before_agent_callback.assert_not_called()
+  spy_run_async_impl.assert_not_called()
+
+  assert len(events) == 1
+  assert events[0].content.parts[0].text == MockPlugin.before_agent_text
+
+
 @pytest.mark.asyncio
 async def test_run_async_with_async_before_agent_callback_noop(
     request: pytest.FixtureRequest,
@@ -486,6 +551,34 @@ async def test_after_agent_callbacks_chain(
         mock_cb.assert_called(expected_calls_count)
 
 
+@pytest.mark.asyncio
+async def test_run_async_after_agent_callback_use_plugin(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+    mock_plugin: MockPlugin,
+):
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      after_agent_callback=_after_agent_callback_noop,
+  )
+  mock_plugin.enable_after_agent_callback = True
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, agent, plugins=[mock_plugin]
+  )
+  spy_after_agent_callback = mocker.spy(agent, 'after_agent_callback')
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_after_agent_callback.assert_not_called()
+  # The first event is regular model response, the second event is
+  # after_agent_callback response.
+  assert len(events) == 2
+  assert events[1].content.parts[0].text == mock_plugin.after_agent_text
+
+
 @pytest.mark.asyncio
 async def test_run_async_after_agent_callback_noop(
     request: pytest.FixtureRequest,
@@ -757,3 +850,7 @@ def test_set_parent_agent_for_sub_agent_twice(
         name=f'{request.function.__name__}_parent_2',
         sub_agents=[sub_agent],
     )
+
+
+if __name__ == '__main__':
+  pytest.main([__file__])

tests/unittests/agents/test_langgraph_agent.py

@@ -17,6 +17,7 @@ from unittest.mock import MagicMock
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.langgraph_agent import LangGraphAgent
 from google.adk.events 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
@@ -169,6 +170,7 @@ async def test_langgraph_agent(
   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",

tests/unittests/agents/test_readonly_context.py

@@ -4,6 +4,7 @@ from unittest.mock import MagicMock
 from google.adk.agents.readonly_context import ReadonlyContext
 import pytest
 
+
 @pytest.fixture
 def mock_invocation_context():
   mock_context = MagicMock()

tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py

@@ -208,9 +208,9 @@ class TestOAuth2CredentialExchanger:
 
     exchanger = OAuth2CredentialExchanger()
 
-    # Mock AUTHLIB_AVIALABLE to False
+    # Mock AUTHLIB_AVAILABLE to False
     with patch(
-        "google.adk.auth.exchanger.oauth2_credential_exchanger.AUTHLIB_AVIALABLE",
+        "google.adk.auth.exchanger.oauth2_credential_exchanger.AUTHLIB_AVAILABLE",
         False,
     ):
       result = await exchanger.exchange(credential, scheme)

tests/unittests/auth/test_auth_handler.py

@@ -456,7 +456,7 @@ class TestExchangeAuthToken:
       self, auth_config_with_auth_code, monkeypatch
   ):
     """Test when token exchange is not supported."""
-    monkeypatch.setattr("google.adk.auth.auth_handler.AUTHLIB_AVIALABLE", False)
+    monkeypatch.setattr("google.adk.auth.auth_handler.AUTHLIB_AVAILABLE", False)
 
     handler = AuthHandler(auth_config_with_auth_code)
     result = await handler.exchange_auth_token()

tests/unittests/cli/utils/test_agent_loader.py

@@ -19,6 +19,7 @@ import tempfile
 from textwrap import dedent
 
 from google.adk.cli.utils.agent_loader import AgentLoader
+from pydantic import ValidationError
 import pytest
 
 
@@ -30,6 +31,8 @@ class TestAgentLoader:
     """Ensure sys.path is restored after each test."""
     original_path = sys.path.copy()
     original_env = os.environ.copy()
+    # Enable WIP features for YAML agent loading tests
+    os.environ["ADK_ALLOW_WIP_FEATURES"] = "true"
     yield
     sys.path[:] = original_path
     # Restore environment variables
@@ -292,7 +295,8 @@ class TestAgentLoader:
       expected_msg_part_1 = "No root_agent found for 'nonexistent_agent'."
       expected_msg_part_2 = (
           "Searched in 'nonexistent_agent.agent.root_agent',"
-          " 'nonexistent_agent.root_agent'."
+          " 'nonexistent_agent.root_agent' and"
+          " 'nonexistent_agent/root_agent.yaml'."
       )
       expected_msg_part_3 = (
           f"Ensure '{agents_dir}/nonexistent_agent' is structured correctly"
@@ -443,3 +447,129 @@ class TestAgentLoader:
       # Now assert path was added
       assert str(temp_path) in sys.path
       assert agent.name == "path_agent"
+
+  def create_yaml_agent_structure(
+      self, temp_dir: Path, agent_name: str, yaml_content: str
+  ):
+    """Create an agent structure with YAML configuration.
+
+    Args:
+        temp_dir: The temporary directory to create the agent in
+        agent_name: Name of the agent
+        yaml_content: YAML content for the root_agent.yaml file
+    """
+    agent_dir = temp_dir / agent_name
+    agent_dir.mkdir()
+
+    # Create root_agent.yaml file
+    yaml_file = agent_dir / "root_agent.yaml"
+    yaml_file.write_text(yaml_content)
+
+  def test_load_agent_from_yaml_config(self):
+    """Test loading an agent from YAML configuration."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+      temp_path = Path(temp_dir)
+      agent_name = "yaml_agent"
+
+      # Create YAML configuration
+      yaml_content = dedent("""
+        agent_class: LlmAgent
+        name: yaml_test_agent
+        model: gemini-2.0-flash
+        instruction: You are a test agent loaded from YAML configuration.
+        description: A test agent created from YAML config
+      """)
+
+      self.create_yaml_agent_structure(temp_path, agent_name, yaml_content)
+
+      # Load the agent
+      loader = AgentLoader(str(temp_path))
+      agent = loader.load_agent(agent_name)
+
+      # Assert agent was loaded correctly
+      assert agent.name == "yaml_test_agent"
+      # Check if it's an LlmAgent before accessing model and instruction
+      from google.adk.agents.llm_agent import LlmAgent
+
+      if isinstance(agent, LlmAgent):
+        assert agent.model == "gemini-2.0-flash"
+        # Handle instruction which can be string or InstructionProvider
+        instruction_text = str(agent.instruction)
+        assert "test agent loaded from YAML" in instruction_text
+
+  def test_yaml_agent_caching_returns_same_instance(self):
+    """Test that loading the same YAML agent twice returns the same instance."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+      temp_path = Path(temp_dir)
+      agent_name = "cached_yaml_agent"
+
+      # Create YAML configuration
+      yaml_content = dedent("""
+        agent_class: LlmAgent
+        name: cached_yaml_test_agent
+        model: gemini-2.0-flash
+        instruction: You are a cached test agent.
+      """)
+
+      self.create_yaml_agent_structure(temp_path, agent_name, yaml_content)
+
+      # Load the agent twice
+      loader = AgentLoader(str(temp_path))
+      agent1 = loader.load_agent(agent_name)
+      agent2 = loader.load_agent(agent_name)
+
+      # Assert same instance is returned
+      assert agent1 is agent2
+      assert agent1.name == agent2.name
+
+  def test_yaml_agent_not_found_error(self):
+    """Test that appropriate error is raised when YAML agent is not found."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+      loader = AgentLoader(temp_dir)
+      agents_dir = temp_dir  # For use in the expected message string
+
+      # Try to load non-existent YAML agent
+      with pytest.raises(ValueError) as exc_info:
+        loader.load_agent("nonexistent_yaml_agent")
+
+      expected_msg_part_1 = "No root_agent found for 'nonexistent_yaml_agent'."
+      expected_msg_part_2 = (
+          "Searched in 'nonexistent_yaml_agent.agent.root_agent',"
+          " 'nonexistent_yaml_agent.root_agent' and"
+          " 'nonexistent_yaml_agent/root_agent.yaml'."
+      )
+      expected_msg_part_3 = (
+          f"Ensure '{agents_dir}/nonexistent_yaml_agent' is structured"
+          " correctly"
+      )
+
+      assert expected_msg_part_1 in str(exc_info.value)
+      assert expected_msg_part_2 in str(exc_info.value)
+      assert expected_msg_part_3 in str(exc_info.value)
+
+  def test_yaml_agent_invalid_yaml_error(self):
+    """Test that appropriate error is raised when YAML is invalid."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+      temp_path = Path(temp_dir)
+      agent_name = "invalid_yaml_agent"
+
+      # Create invalid YAML content with wrong field name
+      invalid_yaml_content = dedent("""
+        agent_type: LlmAgent
+        name: invalid_yaml_test_agent
+        model: gemini-2.0-flash
+        instruction: You are a test agent with invalid YAML
+      """)
+
+      self.create_yaml_agent_structure(
+          temp_path, agent_name, invalid_yaml_content
+      )
+
+      loader = AgentLoader(str(temp_path))
+
+      # Try to load agent with invalid YAML
+      with pytest.raises(ValidationError) as exc_info:
+        loader.load_agent(agent_name)
+
+      # Should raise some form of YAML parsing error
+      assert "Extra inputs are not permitted" in str(exc_info.value)

tests/unittests/evaluation/test_final_response_match_v2.py

@@ -0,0 +1,478 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.evaluation.eval_case import Invocation
+from google.adk.evaluation.eval_metrics import EvalMetric
+from google.adk.evaluation.eval_metrics import JudgeModelOptions
+from google.adk.evaluation.evaluator import EvalStatus
+from google.adk.evaluation.evaluator import PerInvocationResult
+from google.adk.evaluation.final_response_match_v2 import _parse_critique
+from google.adk.evaluation.final_response_match_v2 import FinalResponseMatchV2Evaluator
+from google.adk.evaluation.llm_as_judge_utils import Label
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types as genai_types
+import pytest
+
+
+@pytest.mark.parametrize(
+    "response_text",
+    [
+        """```json
+  {
+    "is_the_agent_response_valid_or_invalid": "valid",
+    "reasoning": "The response is valid."
+  }
+  ```""",
+        """```json
+  {
+    "is_the_agent_response_valid": "undefined label",
+  }
+  ```""",
+    ],
+)
+def test_parse_critique_label_not_found(response_text):
+  label = _parse_critique(response_text)
+  assert label == Label.NOT_FOUND
+
+
+@pytest.mark.parametrize(
+    "response_text",
+    [
+        """```json
+  {
+    "is_the_agent_response_valid": "valid",
+    "reasoning": "The response is valid."
+  }
+  ```""",
+        """```json
+  {
+    "is_the_agent_response_valid": ["valid"],
+    "reasoning": "The response is valid."
+  }
+  ```""",
+        """```json
+  {
+    "is_the_agent_response_valid":\n    [ "valid\n"],
+    "reasoning": "The response is valid."
+  }
+  ```""",
+    ],
+)
+def test_parse_critique(response_text):
+  label = _parse_critique(response_text)
+  assert label == Label.VALID
+
+
+@pytest.mark.parametrize(
+    "response_text",
+    [
+        """```json
+  {
+    "is_the_agent_response_invalid": "invalid",
+    "reasoning": "The response is invalid."
+  }
+  ```""",
+        """```json
+  {
+    "is_the_agent_response_invalid": ["invalid"],
+    "reasoning": "The response is invalid."
+  }
+  ```""",
+        """```json
+  {
+    "is_the_agent_response_invalid":\n    [ "invalid\n"],
+    "reasoning": "The response is invalid."
+  }
+  ```""",
+    ],
+)
+def test_parse_critique_invalid(response_text):
+  label = _parse_critique(response_text)
+  assert label == Label.INVALID
+
+
+def create_test_template() -> str:
+  return """
+This is a test template.
+
+{{
+  "User prompt": {prompt},
+  "Agent response": {response},
+  "Reference response": {golden_response},
+}}
+
+The answer should be a json alone which follows the json structure below:
+{{
+  "is_the_agent_response_valid": [valid or invalid],
+  "reasoning":
+}}
+"""
+
+
+def _create_test_evaluator_gemini(
+    threshold: float,
+) -> FinalResponseMatchV2Evaluator:
+  evaluator = FinalResponseMatchV2Evaluator(
+      EvalMetric(
+          metric_name="final_response_match_v2",
+          threshold=threshold,
+          judge_model_options=JudgeModelOptions(
+              judge_model="gemini-2.5-flash",
+              num_samples=3,
+          ),
+      ),
+  )
+  evaluator._auto_rater_prompt_template = create_test_template()
+  return evaluator
+
+
+def _create_test_invocations(
+    candidate: str, reference: str
+) -> tuple[Invocation, Invocation]:
+  """Returns tuple of (actual_invocation, expected_invocation)."""
+  actual_invocation = Invocation(
+      user_content=genai_types.Content(
+          parts=[genai_types.Part(text="This is a test query.")],
+          role="user",
+      ),
+      final_response=genai_types.Content(
+          parts=[genai_types.Part(text=candidate)],
+          role="model",
+      ),
+  )
+  expected_invocation = Invocation(
+      user_content=genai_types.Content(
+          parts=[genai_types.Part(text="This is a test query.")],
+          role="user",
+      ),
+      final_response=genai_types.Content(
+          parts=[genai_types.Part(text=reference)],
+          role="model",
+      ),
+  )
+  return actual_invocation, expected_invocation
+
+
+def test_format_auto_rater_prompt():
+  evaluator = _create_test_evaluator_gemini(threshold=0.8)
+  actual_invocation, expected_invocation = _create_test_invocations(
+      "candidate text", "reference text"
+  )
+  prompt = evaluator.format_auto_rater_prompt(
+      actual_invocation, expected_invocation
+  )
+  assert prompt == """
+This is a test template.
+
+{
+  "User prompt": This is a test query.,
+  "Agent response": candidate text,
+  "Reference response": reference text,
+}
+
+The answer should be a json alone which follows the json structure below:
+{
+  "is_the_agent_response_valid": [valid or invalid],
+  "reasoning":
+}
+"""
+
+
+def test_convert_auto_rater_response_to_score_valid():
+  evaluator = _create_test_evaluator_gemini(threshold=0.8)
+  auto_rater_response = """```json
+{
+  "is_the_agent_response_valid": "valid",
+  "reasoning": "The response is valid."
+}
+```"""
+  llm_response = LlmResponse(
+      content=genai_types.Content(
+          parts=[genai_types.Part(text=auto_rater_response)],
+          role="model",
+      )
+  )
+  score = evaluator.convert_auto_rater_response_to_score(llm_response)
+  assert score == 1.0
+
+
+def test_convert_auto_rater_response_to_score_invalid():
+  evaluator = _create_test_evaluator_gemini(threshold=0.8)
+  auto_rater_response = """```json
+{
+  "is_the_agent_response_valid": "invalid",
+  "reasoning": "The response is invalid."
+}
+```"""
+  llm_response = LlmResponse(
+      content=genai_types.Content(
+          parts=[genai_types.Part(text=auto_rater_response)],
+          role="model",
+      )
+  )
+  score = evaluator.convert_auto_rater_response_to_score(llm_response)
+  assert score == 0.0
+
+
+def test_convert_auto_rater_response_to_score_invalid_json():
+  evaluator = _create_test_evaluator_gemini(threshold=0.8)
+  llm_response = LlmResponse(
+      content=genai_types.Content(
+          parts=[genai_types.Part(text="invalid json")],
+          role="model",
+      )
+  )
+  score = evaluator.convert_auto_rater_response_to_score(llm_response)
+  assert score is None
+
+
+def test_convert_auto_rater_response_to_score_missing_key():
+  evaluator = _create_test_evaluator_gemini(threshold=0.8)
+  llm_response = LlmResponse(
+      content=genai_types.Content(
+          parts=[genai_types.Part(text="{}")],
+          role="model",
+      )
+  )
+  score = evaluator.convert_auto_rater_response_to_score(llm_response)
+  assert score is None
+
+
+def test_aggregate_per_invocation_samples_none_evaluated():
+  evaluator = _create_test_evaluator_gemini(threshold=0.5)
+
+  actual_invocation, expected_invocation = _create_test_invocations(
+      "candidate text", "reference text"
+  )
+
+  per_invocation_result_samples = [
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=None,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=None,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+  ]
+
+  assert (
+      evaluator.aggregate_per_invocation_samples(per_invocation_result_samples)
+      == per_invocation_result_samples[0]
+  )
+
+
+def test_aggregate_per_invocation_samples_valid():
+  evaluator = _create_test_evaluator_gemini(threshold=0.5)
+
+  actual_invocation, expected_invocation = _create_test_invocations(
+      "candidate text", "reference text"
+  )
+
+  per_invocation_result_samples = [
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=None,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+  ]
+
+  per_invocation_result = evaluator.aggregate_per_invocation_samples(
+      per_invocation_result_samples
+  )
+
+  assert per_invocation_result.score == 1.0
+  assert per_invocation_result.eval_status == EvalStatus.PASSED
+
+
+def test_aggregate_per_invocation_samples_invalid():
+  evaluator = _create_test_evaluator_gemini(threshold=0.5)
+
+  actual_invocation, expected_invocation = _create_test_invocations(
+      "candidate text", "reference text"
+  )
+
+  per_invocation_result_samples = [
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=None,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+  ]
+
+  per_invocation_result = evaluator.aggregate_per_invocation_samples(
+      per_invocation_result_samples
+  )
+
+  assert per_invocation_result.score == 0.0
+  assert per_invocation_result.eval_status == EvalStatus.FAILED
+
+
+def test_aggregate_invocation_results():
+  evaluator = _create_test_evaluator_gemini(threshold=0.5)
+
+  actual_invocation, expected_invocation = _create_test_invocations(
+      "candidate text", "reference text"
+  )
+
+  per_invocation_results = [
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=1.0,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=0.0,
+          eval_status=EvalStatus.FAILED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=None,
+          eval_status=EvalStatus.PASSED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=100.0,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+      PerInvocationResult(
+          actual_invocation=actual_invocation,
+          expected_invocation=expected_invocation,
+          score=None,
+          eval_status=EvalStatus.NOT_EVALUATED,
+      ),
+  ]
+
+  aggregated_result = evaluator.aggregate_invocation_results(
+      per_invocation_results
+  )
+
+  # Only 4 / 8 invocations are evaluated, and 2 / 4 are valid.
+  assert aggregated_result.overall_score == 0.5
+  assert aggregated_result.overall_eval_status == EvalStatus.PASSED

tests/unittests/evaluation/test_llm_as_judge.py

@@ -0,0 +1,221 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Optional
+from unittest.mock import MagicMock
+
+from google.adk.evaluation.eval_case import Invocation
+from google.adk.evaluation.eval_metrics import EvalMetric
+from google.adk.evaluation.eval_metrics import JudgeModelOptions
+from google.adk.evaluation.evaluator import EvalStatus
+from google.adk.evaluation.evaluator import EvaluationResult
+from google.adk.evaluation.evaluator import PerInvocationResult
+from google.adk.evaluation.llm_as_judge import LlmAsJudge
+from google.adk.evaluation.llm_as_judge_utils import get_eval_status
+from google.adk.evaluation.llm_as_judge_utils import get_text_from_content
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types as genai_types
+import pytest
+
+
+class MockLlmAsJudge(LlmAsJudge):
+
+  def format_auto_rater_prompt(
+      self, actual_invocation: Invocation, expected_invocation: Invocation
+  ) -> str:
+    return "formatted prompt"
+
+  def convert_auto_rater_response_to_score(
+      self, llm_response: LlmResponse
+  ) -> Optional[float]:
+    return 1.0
+
+  def aggregate_per_invocation_samples(
+      self,
+      per_invocation_samples: list[PerInvocationResult],
+  ) -> PerInvocationResult:
+    return per_invocation_samples[0]
+
+  def aggregate_invocation_results(
+      self, per_invocation_results: list[PerInvocationResult]
+  ) -> EvaluationResult:
+    return EvaluationResult(
+        overall_score=1.0, overall_eval_status=EvalStatus.PASSED
+    )
+
+
+@pytest.fixture
+def mock_llm_as_judge():
+  return MockLlmAsJudge(
+      EvalMetric(
+          metric_name="test_metric",
+          threshold=0.5,
+          judge_model_options=JudgeModelOptions(
+              judge_model="gemini-2.5-flash",
+              judge_model_config=genai_types.GenerateContentConfig(),
+              num_samples=3,
+          ),
+      ),
+  )
+
+
+def test_get_text_from_content():
+  content = genai_types.Content(
+      parts=[
+          genai_types.Part(text="This is a test text."),
+          genai_types.Part(text="This is another test text."),
+      ],
+      role="model",
+  )
+  assert (
+      get_text_from_content(content)
+      == "This is a test text.\nThis is another test text."
+  )
+
+
+def test_get_eval_status():
+  assert get_eval_status(score=0.8, threshold=0.8) == EvalStatus.PASSED
+  assert get_eval_status(score=0.7, threshold=0.8) == EvalStatus.FAILED
+  assert get_eval_status(score=0.8, threshold=0.9) == EvalStatus.FAILED
+  assert get_eval_status(score=0.9, threshold=0.8) == EvalStatus.PASSED
+  assert get_eval_status(score=None, threshold=0.8) == EvalStatus.NOT_EVALUATED
+
+
+def test_llm_as_judge_init_missing_judge_model_options():
+  with pytest.raises(ValueError):
+    MockLlmAsJudge(
+        EvalMetric(metric_name="test_metric", threshold=0.8),
+    )
+
+
+def test_llm_as_judge_init_unregistered_model():
+  with pytest.raises(ValueError):
+    MockLlmAsJudge(
+        EvalMetric(
+            metric_name="test_metric",
+            threshold=0.8,
+            judge_model_options=JudgeModelOptions(
+                judge_model="unregistered_model",
+            ),
+        ),
+    )
+
+
+@pytest.fixture
+def mock_judge_model():
+  mock_judge_model = MagicMock()
+
+  async def mock_generate_content_async(llm_request):
+    yield LlmResponse(
+        content=genai_types.Content(
+            parts=[genai_types.Part(text="auto rater response")],
+        )
+    )
+
+  mock_judge_model.generate_content_async = mock_generate_content_async
+  return mock_judge_model
+
+
+@pytest.mark.asyncio
+async def test_evaluate_invocations_with_mock(
+    mock_llm_as_judge, mock_judge_model
+):
+  mock_llm_as_judge._judge_model = mock_judge_model
+
+  mock_format_auto_rater_prompt = MagicMock(
+      wraps=mock_llm_as_judge.format_auto_rater_prompt
+  )
+  mock_llm_as_judge.format_auto_rater_prompt = mock_format_auto_rater_prompt
+
+  mock_convert_auto_rater_response_to_score = MagicMock(
+      wraps=mock_llm_as_judge.convert_auto_rater_response_to_score
+  )
+  mock_llm_as_judge.convert_auto_rater_response_to_score = (
+      mock_convert_auto_rater_response_to_score
+  )
+
+  mock_aggregate_per_invocation_samples = MagicMock(
+      wraps=mock_llm_as_judge.aggregate_per_invocation_samples
+  )
+  mock_llm_as_judge.aggregate_per_invocation_samples = (
+      mock_aggregate_per_invocation_samples
+  )
+
+  mock_aggregate_invocation_results = MagicMock(
+      wraps=mock_llm_as_judge.aggregate_invocation_results
+  )
+  mock_llm_as_judge.aggregate_invocation_results = (
+      mock_aggregate_invocation_results
+  )
+
+  actual_invocations = [
+      Invocation(
+          invocation_id="id1",
+          user_content=genai_types.Content(
+              parts=[genai_types.Part(text="user content 1")],
+              role="user",
+          ),
+          final_response=genai_types.Content(
+              parts=[genai_types.Part(text="final response 1")],
+              role="model",
+          ),
+      ),
+      Invocation(
+          invocation_id="id2",
+          user_content=genai_types.Content(
+              parts=[genai_types.Part(text="user content 2")],
+              role="user",
+          ),
+          final_response=genai_types.Content(
+              parts=[genai_types.Part(text="final response 2")],
+              role="model",
+          ),
+      ),
+  ]
+  expected_invocations = [
+      Invocation(
+          invocation_id="id1",
+          user_content=genai_types.Content(
+              parts=[genai_types.Part(text="user content 1")],
+              role="user",
+          ),
+          final_response=genai_types.Content(
+              parts=[genai_types.Part(text="expected response 1")],
+              role="model",
+          ),
+      ),
+      Invocation(
+          invocation_id="id2",
+          user_content=genai_types.Content(
+              parts=[genai_types.Part(text="user content 2")],
+              role="user",
+          ),
+          final_response=genai_types.Content(
+              parts=[genai_types.Part(text="expected response 2")],
+              role="model",
+          ),
+      ),
+  ]
+
+  result = await mock_llm_as_judge.evaluate_invocations(
+      actual_invocations, expected_invocations
+  )
+
+  # Assertions
+  assert result.overall_score == 1.0
+  assert mock_llm_as_judge.format_auto_rater_prompt.call_count == 2
+  assert mock_llm_as_judge.convert_auto_rater_response_to_score.call_count == 6
+  assert mock_llm_as_judge.aggregate_invocation_results.call_count == 1

tests/unittests/evaluation/test_local_eval_service.py

@@ -0,0 +1,144 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.errors.not_found_error import NotFoundError
+from google.adk.evaluation.base_eval_service import InferenceConfig
+from google.adk.evaluation.base_eval_service import InferenceRequest
+from google.adk.evaluation.eval_set import EvalCase
+from google.adk.evaluation.eval_set import EvalSet
+from google.adk.evaluation.eval_sets_manager import EvalSetsManager
+from google.adk.evaluation.local_eval_service import LocalEvalService
+from google.adk.models.registry import LLMRegistry
+import pytest
+
+
+@pytest.fixture
+def mock_eval_sets_manager():
+  return mock.create_autospec(EvalSetsManager)
+
+
+@pytest.fixture
+def dummy_agent():
+  llm = LLMRegistry.new_llm("gemini-pro")
+  return LlmAgent(name="test_agent", model=llm)
+
+
+@pytest.fixture
+def eval_service(dummy_agent, mock_eval_sets_manager):
+  return LocalEvalService(
+      root_agent=dummy_agent,
+      eval_sets_manager=mock_eval_sets_manager,
+  )
+
+
+@pytest.mark.asyncio
+async def test_perform_inference_success(
+    eval_service, dummy_agent, mock_eval_sets_manager
+):
+  eval_set = EvalSet(
+      eval_set_id="test_eval_set",
+      eval_cases=[
+          EvalCase(eval_id="case1", conversation=[], session_input=None),
+          EvalCase(eval_id="case2", conversation=[], session_input=None),
+      ],
+  )
+  mock_eval_sets_manager.get_eval_set.return_value = eval_set
+
+  mock_inference_result = mock.MagicMock()
+  eval_service._perform_inference_sigle_eval_item = mock.AsyncMock(
+      return_value=mock_inference_result
+  )
+
+  inference_request = InferenceRequest(
+      app_name="test_app",
+      eval_set_id="test_eval_set",
+      inference_config=InferenceConfig(parallelism=2),
+  )
+
+  results = []
+  async for result in eval_service.perform_inference(inference_request):
+    results.append(result)
+
+  assert len(results) == 2
+  assert results[0] == mock_inference_result
+  assert results[1] == mock_inference_result
+  mock_eval_sets_manager.get_eval_set.assert_called_once_with(
+      app_name="test_app", eval_set_id="test_eval_set"
+  )
+  assert eval_service._perform_inference_sigle_eval_item.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_perform_inference_with_case_ids(
+    eval_service, dummy_agent, mock_eval_sets_manager
+):
+  eval_set = EvalSet(
+      eval_set_id="test_eval_set",
+      eval_cases=[
+          EvalCase(eval_id="case1", conversation=[], session_input=None),
+          EvalCase(eval_id="case2", conversation=[], session_input=None),
+          EvalCase(eval_id="case3", conversation=[], session_input=None),
+      ],
+  )
+  mock_eval_sets_manager.get_eval_set.return_value = eval_set
+
+  mock_inference_result = mock.MagicMock()
+  eval_service._perform_inference_sigle_eval_item = mock.AsyncMock(
+      return_value=mock_inference_result
+  )
+
+  inference_request = InferenceRequest(
+      app_name="test_app",
+      eval_set_id="test_eval_set",
+      eval_case_ids=["case1", "case3"],
+      inference_config=InferenceConfig(parallelism=1),
+  )
+
+  results = []
+  async for result in eval_service.perform_inference(inference_request):
+    results.append(result)
+
+  assert len(results) == 2
+  eval_service._perform_inference_sigle_eval_item.assert_any_call(
+      app_name="test_app",
+      eval_set_id="test_eval_set",
+      eval_case=eval_set.eval_cases[0],
+      root_agent=dummy_agent,
+  )
+  eval_service._perform_inference_sigle_eval_item.assert_any_call(
+      app_name="test_app",
+      eval_set_id="test_eval_set",
+      eval_case=eval_set.eval_cases[2],
+      root_agent=dummy_agent,
+  )
+
+
+@pytest.mark.asyncio
+async def test_perform_inference_eval_set_not_found(
+    eval_service, mock_eval_sets_manager
+):
+  mock_eval_sets_manager.get_eval_set.return_value = None
+
+  inference_request = InferenceRequest(
+      app_name="test_app",
+      eval_set_id="not_found_set",
+      inference_config=InferenceConfig(parallelism=1),
+  )
+
+  with pytest.raises(NotFoundError):
+    async for _ in eval_service.perform_inference(inference_request):
+      pass

tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py

@@ -0,0 +1,164 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk.agents import Agent
+from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class BaseLlmFlowForTesting(BaseLlmFlow):
+  """Test implementation of BaseLlmFlow for testing purposes."""
+
+  pass
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_partial_event():
+  """Test that run_async breaks when the last event is partial."""
+  # Create a mock model that returns partial responses
+  partial_response = LlmResponse(
+      content=types.Content(
+          role='model', parts=[types.Part.from_text(text='Partial response')]
+      ),
+      partial=True,
+  )
+
+  mock_model = testing_utils.MockModel.create(responses=[partial_response])
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have one event (the partial response)
+  assert len(events) == 1
+  assert events[0].partial is True
+  assert events[0].content.parts[0].text == 'Partial response'
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_final_response():
+  """Test that run_async breaks when the last event is a final response."""
+  # Create a mock model that returns a final response
+  final_response = LlmResponse(
+      content=types.Content(
+          role='model', parts=[types.Part.from_text(text='Final response')]
+      ),
+      partial=False,
+      error_code=types.FinishReason.STOP,
+  )
+
+  mock_model = testing_utils.MockModel.create(responses=[final_response])
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have one event (the final response)
+  assert len(events) == 1
+  assert events[0].partial is False
+  assert events[0].content.parts[0].text == 'Final response'
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_no_last_event():
+  """Test that run_async breaks when there is no last event."""
+  # Create a mock model that returns an empty response (no content)
+  empty_response = LlmResponse(content=None, partial=False)
+
+  mock_model = testing_utils.MockModel.create(responses=[empty_response])
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have no events because empty responses are filtered out
+  assert len(events) == 0
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_first_partial_response():
+  """Test run_async breaks on the first partial response."""
+  # Create responses with mixed partial states
+  partial_response = LlmResponse(
+      content=types.Content(
+          role='model', parts=[types.Part.from_text(text='Partial response')]
+      ),
+      partial=True,
+  )
+
+  # These won't be reached because the flow breaks on the first partial
+  non_partial_response = LlmResponse(
+      content=types.Content(
+          role='model',
+          parts=[types.Part.from_text(text='Non-partial response')],
+      ),
+      partial=False,
+  )
+
+  final_partial_response = LlmResponse(
+      content=types.Content(
+          role='model',
+          parts=[types.Part.from_text(text='Final partial response')],
+      ),
+      partial=True,
+  )
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[partial_response, non_partial_response, final_partial_response]
+  )
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have only one event, breaking on the first partial response
+  assert len(events) == 1
+  assert events[0].partial is True
+  assert events[0].content.parts[0].text == 'Partial response'

tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py

@@ -0,0 +1,128 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from google.adk.agents import Agent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.models import LlmRequest
+from google.adk.models import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class MockPlugin(BasePlugin):
+  before_model_text = 'before_model_text from MockPlugin'
+  after_model_text = 'after_model_text from MockPlugin'
+
+  def __init__(self, name='mock_plugin'):
+    self.name = name
+    self.enable_before_model_callback = False
+    self.enable_after_model_callback = False
+    self.before_model_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.before_model_text)]
+        )
+    )
+    self.after_model_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.after_model_text)]
+        )
+    )
+
+  async def before_model_callback(
+      self, *, callback_context: CallbackContext, llm_request: LlmRequest
+  ) -> Optional[LlmResponse]:
+    if not self.enable_before_model_callback:
+      return None
+    return self.before_model_response
+
+  async def after_model_callback(
+      self, *, callback_context: CallbackContext, llm_response: LlmResponse
+  ) -> Optional[LlmResponse]:
+    if not self.enable_after_model_callback:
+      return None
+    return self.after_model_response
+
+
+CANONICAL_MODEL_CALLBACK_CONTENT = 'canonical_model_callback_content'
+
+
+def canonical_agent_model_callback(**kwargs) -> Optional[LlmResponse]:
+  return LlmResponse(
+      content=testing_utils.ModelContent(
+          [types.Part.from_text(text=CANONICAL_MODEL_CALLBACK_CONTENT)]
+      )
+  )
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
+def test_before_model_callback_with_plugin(mock_plugin):
+  """Tests that the model response is overridden by before_model_callback from the plugin."""
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  mock_plugin.enable_before_model_callback = True
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', mock_plugin.before_model_text),
+  ]
+
+
+def test_before_model_fallback_canonical_callback(mock_plugin):
+  """Tests that when plugin returns empty response, the model response is overridden by the canonical agent model callback."""
+  responses = ['model_response']
+  mock_plugin.enable_before_model_callback = False
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=canonical_agent_model_callback,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', CANONICAL_MODEL_CALLBACK_CONTENT),
+  ]
+
+
+def test_before_model_callback_fallback_model(mock_plugin):
+  """Tests that the model response is executed normally when both plugin and canonical agent model callback return empty response."""
+  responses = ['model_response']
+  mock_plugin.enable_before_model_callback = False
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+if __name__ == '__main__':
+  pytest.main([__file__])

tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py

@@ -0,0 +1,128 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from google.adk.agents import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import handle_function_calls_async
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class MockPlugin(BasePlugin):
+  before_tool_response = {"MockPlugin": "before_tool_response from MockPlugin"}
+  after_tool_response = {"MockPlugin": "after_tool_response from MockPlugin"}
+
+  def __init__(self, name="mock_plugin"):
+    self.name = name
+    self.enable_before_tool_callback = False
+    self.enable_after_tool_callback = False
+
+  async def before_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Optional[dict]:
+    if not self.enable_before_tool_callback:
+      return None
+    return self.before_tool_response
+
+  async def after_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      result: dict,
+  ) -> Optional[dict]:
+    if not self.enable_after_tool_callback:
+      return None
+    return self.after_tool_response
+
+
+@pytest.fixture
+def mock_tool():
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  return FunctionTool(simple_fn)
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
+async def invoke_tool_with_plugin(mock_tool, mock_plugin) -> Optional[Event]:
+  """Invokes a tool with a plugin."""
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[mock_tool],
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content="", plugins=[mock_plugin]
+  )
+  # Build function call event
+  function_call = types.FunctionCall(name=mock_tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {mock_tool.name: mock_tool}
+  return await handle_function_calls_async(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+@pytest.mark.asyncio
+async def test_async_before_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_before_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.before_tool_response
+
+
+@pytest.mark.asyncio
+async def test_async_after_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_after_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.after_tool_response
+
+
+if __name__ == "__main__":
+  pytest.main([__file__])

tests/unittests/plugins/test_base_plugin.py

@@ -0,0 +1,239 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from unittest.mock import Mock
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+
+class TestablePlugin(BasePlugin):
+  __test__ = False
+  """A concrete implementation of BasePlugin for testing purposes."""
+  pass
+
+
+class FullOverridePlugin(BasePlugin):
+  __test__ = False
+
+  """A plugin that overrides every single callback method for testing."""
+
+  def __init__(self, name: str = "full_override"):
+    super().__init__(name)
+
+  async def on_user_message_callback(self, **kwargs) -> str:
+    return "overridden_on_user_message"
+
+  async def before_run_callback(self, **kwargs) -> str:
+    return "overridden_before_run"
+
+  async def after_run_callback(self, **kwargs) -> str:
+    return "overridden_after_run"
+
+  async def on_event_callback(self, **kwargs) -> str:
+    return "overridden_on_event"
+
+  async def before_agent_callback(self, **kwargs) -> str:
+    return "overridden_before_agent"
+
+  async def after_agent_callback(self, **kwargs) -> str:
+    return "overridden_after_agent"
+
+  async def before_tool_callback(self, **kwargs) -> str:
+    return "overridden_before_tool"
+
+  async def after_tool_callback(self, **kwargs) -> str:
+    return "overridden_after_tool"
+
+  async def before_model_callback(self, **kwargs) -> str:
+    return "overridden_before_model"
+
+  async def after_model_callback(self, **kwargs) -> str:
+    return "overridden_after_model"
+
+
+def test_base_plugin_initialization():
+  """Tests that a plugin is initialized with the correct name."""
+  plugin_name = "my_test_plugin"
+  plugin = TestablePlugin(name=plugin_name)
+  assert plugin.name == plugin_name
+
+
+@pytest.mark.asyncio
+async def test_base_plugin_default_callbacks_return_none():
+  """Tests that the default (non-overridden) callbacks in BasePlugin exist
+
+  and return None as expected.
+  """
+  plugin = TestablePlugin(name="default_plugin")
+
+  # Mocking all necessary context objects
+  mock_context = Mock()
+  mock_user_message = Mock()
+
+  # The default implementations should do nothing and return None.
+  assert (
+      await plugin.on_user_message_callback(
+          user_message=mock_user_message,
+          invocation_context=mock_context,
+      )
+      is None
+  )
+  assert (
+      await plugin.before_run_callback(invocation_context=mock_context) is None
+  )
+  assert (
+      await plugin.after_run_callback(invocation_context=mock_context) is None
+  )
+  assert (
+      await plugin.on_event_callback(
+          invocation_context=mock_context, event=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.before_agent_callback(
+          agent=mock_context, callback_context=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.after_agent_callback(
+          agent=mock_context, callback_context=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.before_tool_callback(
+          tool=mock_context, tool_args={}, tool_context=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.after_tool_callback(
+          tool=mock_context, tool_args={}, tool_context=mock_context, result={}
+      )
+      is None
+  )
+  assert (
+      await plugin.before_model_callback(
+          callback_context=mock_context, llm_request=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.after_model_callback(
+          callback_context=mock_context, llm_response=mock_context
+      )
+      is None
+  )
+
+
+@pytest.mark.asyncio
+async def test_base_plugin_all_callbacks_can_be_overridden():
+  """Verifies that a user can create a subclass of BasePlugin and that all
+
+  overridden methods are correctly called.
+  """
+  plugin = FullOverridePlugin()
+
+  # Create mock objects for all required arguments. We don't need real
+  # objects, just placeholders to satisfy the method signatures.
+  mock_user_message = Mock(spec=types.Content)
+  mock_invocation_context = Mock(spec=InvocationContext)
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_agent = Mock(spec=BaseAgent)
+  mock_tool = Mock(spec=BaseTool)
+  mock_tool_context = Mock(spec=ToolContext)
+  mock_llm_request = Mock(spec=LlmRequest)
+  mock_llm_response = Mock(spec=LlmResponse)
+  mock_event = Mock(spec=Event)
+
+  # Call each method and assert it returns the unique string from the override.
+  # This proves that the subclass's method was executed.
+  assert (
+      await plugin.on_user_message_callback(
+          user_message=mock_user_message,
+          invocation_context=mock_invocation_context,
+      )
+      == "overridden_on_user_message"
+  )
+  assert (
+      await plugin.before_run_callback(
+          invocation_context=mock_invocation_context
+      )
+      == "overridden_before_run"
+  )
+  assert (
+      await plugin.after_run_callback(
+          invocation_context=mock_invocation_context
+      )
+      == "overridden_after_run"
+  )
+  assert (
+      await plugin.on_event_callback(
+          invocation_context=mock_invocation_context, event=mock_event
+      )
+      == "overridden_on_event"
+  )
+  assert (
+      await plugin.before_agent_callback(
+          agent=mock_agent, callback_context=mock_callback_context
+      )
+      == "overridden_before_agent"
+  )
+  assert (
+      await plugin.after_agent_callback(
+          agent=mock_agent, callback_context=mock_callback_context
+      )
+      == "overridden_after_agent"
+  )
+  assert (
+      await plugin.before_model_callback(
+          callback_context=mock_callback_context, llm_request=mock_llm_request
+      )
+      == "overridden_before_model"
+  )
+  assert (
+      await plugin.after_model_callback(
+          callback_context=mock_callback_context, llm_response=mock_llm_response
+      )
+      == "overridden_after_model"
+  )
+  assert (
+      await plugin.before_tool_callback(
+          tool=mock_tool, tool_args={}, tool_context=mock_tool_context
+      )
+      == "overridden_before_tool"
+  )
+  assert (
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args={},
+          tool_context=mock_tool_context,
+          result={},
+      )
+      == "overridden_after_tool"
+  )

tests/unittests/plugins/test_plugin_manager.py

@@ -0,0 +1,250 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the PluginManager."""
+
+from __future__ import annotations
+
+from unittest.mock import Mock
+
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+# Assume the following path to your modules
+# You might need to adjust this based on your project structure.
+from google.adk.plugins.plugin_manager import PluginCallbackName
+from google.adk.plugins.plugin_manager import PluginManager
+import pytest
+
+
+# A helper class to use in tests instead of mocks.
+# This makes tests more explicit and easier to debug.
+class TestPlugin(BasePlugin):
+  __test__ = False
+  """
+  A test plugin that can be configured to return specific values or raise
+  exceptions for any callback, and it logs which callbacks were invoked.
+  """
+
+  def __init__(self, name: str):
+    super().__init__(name)
+    # A log to track the names of callbacks that have been called.
+    self.call_log: list[PluginCallbackName] = []
+    # A map to configure return values for specific callbacks.
+    self.return_values: dict[PluginCallbackName, any] = {}
+    # A map to configure exceptions to be raised by specific callbacks.
+    self.exceptions_to_raise: dict[PluginCallbackName, Exception] = {}
+
+  async def _handle_callback(self, name: PluginCallbackName):
+    """Generic handler for all callback methods."""
+    self.call_log.append(name)
+    if name in self.exceptions_to_raise:
+      raise self.exceptions_to_raise[name]
+    return self.return_values.get(name)
+
+  # Implement all callback methods from the BasePlugin interface.
+  async def on_user_message_callback(self, **kwargs):
+    return await self._handle_callback("on_user_message_callback")
+
+  async def before_run_callback(self, **kwargs):
+    return await self._handle_callback("before_run_callback")
+
+  async def after_run_callback(self, **kwargs):
+    return await self._handle_callback("after_run_callback")
+
+  async def on_event_callback(self, **kwargs):
+    return await self._handle_callback("on_event_callback")
+
+  async def before_agent_callback(self, **kwargs):
+    return await self._handle_callback("before_agent_callback")
+
+  async def after_agent_callback(self, **kwargs):
+    return await self._handle_callback("after_agent_callback")
+
+  async def before_tool_callback(self, **kwargs):
+    return await self._handle_callback("before_tool_callback")
+
+  async def after_tool_callback(self, **kwargs):
+    return await self._handle_callback("after_tool_callback")
+
+  async def before_model_callback(self, **kwargs):
+    return await self._handle_callback("before_model_callback")
+
+  async def after_model_callback(self, **kwargs):
+    return await self._handle_callback("after_model_callback")
+
+
+@pytest.fixture
+def service() -> PluginManager:
+  """Provides a clean PluginManager instance for each test."""
+  return PluginManager()
+
+
+@pytest.fixture
+def plugin1() -> TestPlugin:
+  """Provides a clean instance of our test plugin named 'plugin1'."""
+  return TestPlugin(name="plugin1")
+
+
+@pytest.fixture
+def plugin2() -> TestPlugin:
+  """Provides a clean instance of our test plugin named 'plugin2'."""
+  return TestPlugin(name="plugin2")
+
+
+def test_register_and_get_plugin(service: PluginManager, plugin1: TestPlugin):
+  """Tests successful registration and retrieval of a plugin."""
+  service.register_plugin(plugin1)
+
+  assert len(service.plugins) == 1
+  assert service.plugins[0] is plugin1
+  assert service.get_plugin("plugin1") is plugin1
+
+
+def test_register_duplicate_plugin_name_raises_value_error(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that registering a plugin with a duplicate name raises an error."""
+  plugin1_duplicate = TestPlugin(name="plugin1")
+  service.register_plugin(plugin1)
+
+  with pytest.raises(
+      ValueError, match="Plugin with name 'plugin1' already registered."
+  ):
+    service.register_plugin(plugin1_duplicate)
+
+
+@pytest.mark.asyncio
+async def test_early_exit_stops_subsequent_plugins(
+    service: PluginManager, plugin1: TestPlugin, plugin2: TestPlugin
+):
+  """Tests the core "early exit" logic: if a plugin returns a value,
+
+  subsequent plugins for that callback should not be executed.
+  """
+  # Configure plugin1 to return a value, simulating a cache hit.
+  mock_response = Mock(spec=LlmResponse)
+  plugin1.return_values["before_run_callback"] = mock_response
+
+  service.register_plugin(plugin1)
+  service.register_plugin(plugin2)
+
+  # Execute the callback chain.
+  result = await service.run_before_run_callback(invocation_context=Mock())
+
+  # Assert that the final result is the one returned by the first plugin.
+  assert result is mock_response
+  # Assert that the first plugin was called.
+  assert "before_run_callback" in plugin1.call_log
+  # CRITICAL: Assert that the second plugin was never called.
+  assert "before_run_callback" not in plugin2.call_log
+
+
+@pytest.mark.asyncio
+async def test_normal_flow_all_plugins_are_called(
+    service: PluginManager, plugin1: TestPlugin, plugin2: TestPlugin
+):
+  """Tests that if no plugin returns a value, all plugins in the chain
+
+  are executed in order.
+  """
+  # By default, plugins are configured to return None.
+  service.register_plugin(plugin1)
+  service.register_plugin(plugin2)
+
+  result = await service.run_before_run_callback(invocation_context=Mock())
+
+  # The final result should be None as no plugin interrupted the flow.
+  assert result is None
+  # Both plugins must have been called.
+  assert "before_run_callback" in plugin1.call_log
+  assert "before_run_callback" in plugin2.call_log
+
+
+@pytest.mark.asyncio
+async def test_plugin_exception_is_wrapped_in_runtime_error(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that if a plugin callback raises an exception, the PluginManager
+
+  catches it and raises a descriptive RuntimeError.
+  """
+  # Configure the plugin to raise an error during a specific callback.
+  original_exception = ValueError("Something went wrong inside the plugin!")
+  plugin1.exceptions_to_raise["before_run_callback"] = original_exception
+  service.register_plugin(plugin1)
+
+  with pytest.raises(RuntimeError) as excinfo:
+    await service.run_before_run_callback(invocation_context=Mock())
+
+  # Check that the error message is informative.
+  assert "Error in plugin 'plugin1'" in str(excinfo.value)
+  assert "before_run_callback" in str(excinfo.value)
+  # Check that the original exception is chained for better tracebacks.
+  assert excinfo.value.__cause__ is original_exception
+
+
+@pytest.mark.asyncio
+async def test_all_callbacks_are_supported(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that all callbacks defined in the BasePlugin interface are supported
+
+  by the PluginManager.
+  """
+  service.register_plugin(plugin1)
+  mock_context = Mock()
+  mock_user_message = Mock()
+
+  # Test all callbacks
+  await service.run_on_user_message_callback(
+      user_message=mock_user_message, invocation_context=mock_context
+  )
+  await service.run_before_run_callback(invocation_context=mock_context)
+  await service.run_after_run_callback(invocation_context=mock_context)
+  await service.run_on_event_callback(
+      invocation_context=mock_context, event=mock_context
+  )
+  await service.run_before_agent_callback(
+      agent=mock_context, callback_context=mock_context
+  )
+  await service.run_after_agent_callback(
+      agent=mock_context, callback_context=mock_context
+  )
+  await service.run_before_tool_callback(
+      tool=mock_context, tool_args={}, tool_context=mock_context
+  )
+  await service.run_after_tool_callback(
+      tool=mock_context, tool_args={}, tool_context=mock_context, result={}
+  )
+  await service.run_before_model_callback(
+      callback_context=mock_context, llm_request=mock_context
+  )
+  await service.run_after_model_callback(
+      callback_context=mock_context, llm_response=mock_context
+  )
+
+  # Verify all callbacks were logged
+  expected_callbacks = [
+      "on_user_message_callback",
+      "before_run_callback",
+      "after_run_callback",
+      "on_event_callback",
+      "before_agent_callback",
+      "after_agent_callback",
+      "before_tool_callback",
+      "after_tool_callback",
+      "before_model_callback",
+      "after_model_callback",
+  ]
+  assert set(plugin1.call_log) == set(expected_callbacks)

tests/unittests/test_runners.py

@@ -15,13 +15,20 @@
 from typing import Optional
 
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
 from google.adk.events.event import Event
+from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
 from google.genai import types
+import pytest
+
+TEST_APP_ID = "test_app"
+TEST_USER_ID = "test_user"
+TEST_SESSION_ID = "test_session"
 
 
 class MockAgent(BaseAgent):
@@ -72,6 +79,51 @@ class MockLlmAgent(LlmAgent):
     )
 
 
+class MockPlugin(BasePlugin):
+  """Mock plugin for unit testing."""
+
+  ON_USER_CALLBACK_MSG = (
+      "Modified user message ON_USER_CALLBACK_MSG from MockPlugin"
+  )
+  ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin"
+
+  def __init__(self):
+    super().__init__(name="mock_plugin")
+    self.enable_user_message_callback = False
+    self.enable_event_callback = False
+
+  async def on_user_message_callback(
+      self,
+      *,
+      invocation_context: InvocationContext,
+      user_message: types.Content,
+  ) -> Optional[types.Content]:
+    if not self.enable_user_message_callback:
+      return None
+    return types.Content(
+        role="model",
+        parts=[types.Part(text=self.ON_USER_CALLBACK_MSG)],
+    )
+
+  async def on_event_callback(
+      self, *, invocation_context: InvocationContext, event: Event
+  ) -> Optional[Event]:
+    if not self.enable_event_callback:
+      return None
+    return Event(
+        invocation_id="",
+        author="",
+        content=types.Content(
+            parts=[
+                types.Part(
+                    text=self.ON_EVENT_CALLBACK_MSG,
+                )
+            ],
+            role=event.content.role,
+        ),
+    )
+
+
 class TestRunnerFindAgentToRun:
   """Tests for Runner._find_agent_to_run method."""
 
@@ -308,3 +360,73 @@ class TestRunnerFindAgentToRun:
     # MockAgent inherits from BaseAgent, not LlmAgent, so it should return False
     result = self.runner._is_transferable_across_agent_tree(non_llm_agent)
     assert result is False
+
+
+class TestRunnerWithPlugins:
+  """Tests for Runner with plugins."""
+
+  def setup_method(self):
+    self.plugin = MockPlugin()
+    self.session_service = InMemorySessionService()
+    self.artifact_service = InMemoryArtifactService()
+    self.root_agent = MockLlmAgent("root_agent")
+    self.runner = Runner(
+        app_name="test_app",
+        agent=MockLlmAgent("test_agent"),
+        session_service=self.session_service,
+        artifact_service=self.artifact_service,
+        plugins=[self.plugin],
+    )
+
+  async def run_test(self, original_user_input="Hello") -> list[Event]:
+    """Prepares the test by creating a session and running the runner."""
+    await self.session_service.create_session(
+        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+    )
+    events = []
+    async for event in self.runner.run_async(
+        user_id=TEST_USER_ID,
+        session_id=TEST_SESSION_ID,
+        new_message=types.Content(
+            role="user", parts=[types.Part(text=original_user_input)]
+        ),
+    ):
+      events.append(event)
+    return events
+
+  @pytest.mark.asyncio
+  async def test_runner_is_initialized_with_plugins(self):
+    """Test that the runner is initialized with plugins."""
+    await self.run_test()
+
+    assert self.runner.plugin_manager is not None
+
+  @pytest.mark.asyncio
+  async def test_runner_modifies_user_message_before_execution(self):
+    """Test that the runner modifies the user message before execution."""
+    original_user_input = "original_input"
+    self.plugin.enable_user_message_callback = True
+
+    await self.run_test(original_user_input=original_user_input)
+    session = await self.session_service.get_session(
+        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+    )
+    generated_event = session.events[0]
+    modified_user_message = generated_event.content.parts[0].text
+
+    assert modified_user_message == MockPlugin.ON_USER_CALLBACK_MSG
+
+  @pytest.mark.asyncio
+  async def test_runner_modifies_event_after_execution(self):
+    """Test that the runner modifies the event after execution."""
+    self.plugin.enable_event_callback = True
+
+    events = await self.run_test()
+    generated_event = events[0]
+    modified_event_message = generated_event.content.parts[0].text
+
+    assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG
+
+
+if __name__ == "__main__":
+  pytest.main([__file__])

tests/unittests/testing_utils.py

@@ -30,6 +30,8 @@ from google.adk.models.base_llm import BaseLlm
 from google.adk.models.base_llm_connection import BaseLlmConnection
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.plugins.plugin_manager import PluginManager
 from google.adk.runners import InMemoryRunner as AfInMemoryRunner
 from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
@@ -57,7 +59,10 @@ class ModelContent(types.Content):
 
 
 async def create_invocation_context(
-    agent: Agent, user_content: str = '', run_config: RunConfig = None
+    agent: Agent,
+    user_content: str = '',
+    run_config: RunConfig = None,
+    plugins: list[BasePlugin] = [],
 ):
   invocation_id = 'test_id'
   artifact_service = InMemoryArtifactService()
@@ -67,6 +72,7 @@ async def create_invocation_context(
       artifact_service=artifact_service,
       session_service=session_service,
       memory_service=memory_service,
+      plugin_manager=PluginManager(plugins=plugins),
       invocation_id=invocation_id,
       agent=agent,
       session=await session_service.create_session(
@@ -165,6 +171,7 @@ class InMemoryRunner:
       self,
       root_agent: Union[Agent, LlmAgent],
       response_modalities: list[str] = None,
+      plugins: list[BasePlugin] = [],
   ):
     self.root_agent = root_agent
     self.runner = Runner(
@@ -173,6 +180,7 @@ class InMemoryRunner:
         artifact_service=InMemoryArtifactService(),
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
+        plugins=plugins,
     )
     self.session_id = None
 

tests/unittests/tools/google_api_tool/test_google_api_tool.py

@@ -0,0 +1,145 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import ServiceAccount
+from google.adk.auth.auth_credential import ServiceAccountCredential
+from google.adk.tools.google_api_tool.google_api_tool import GoogleApiTool
+from google.adk.tools.openapi_tool import RestApiTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai.types import FunctionDeclaration
+import pytest
+
+
+@pytest.fixture
+def mock_rest_api_tool():
+  """Fixture for a mock RestApiTool."""
+  mock_tool = mock.MagicMock(spec=RestApiTool)
+  mock_tool.name = "test_tool"
+  mock_tool.description = "Test Tool Description"
+  mock_tool.is_long_running = False
+  mock_tool._get_declaration.return_value = FunctionDeclaration(
+      name="test_function", description="Test function description"
+  )
+  mock_tool.run_async.return_value = {"result": "success"}
+  return mock_tool
+
+
+@pytest.fixture
+def mock_tool_context():
+  """Fixture for a mock ToolContext."""
+  return mock.MagicMock(spec=ToolContext)
+
+
+class TestGoogleApiTool:
+  """Test suite for the GoogleApiTool class."""
+
+  def test_init(self, mock_rest_api_tool):
+    """Test GoogleApiTool initialization."""
+    tool = GoogleApiTool(mock_rest_api_tool)
+
+    assert tool.name == "test_tool"
+    assert tool.description == "Test Tool Description"
+    assert tool.is_long_running is False
+    assert tool._rest_api_tool == mock_rest_api_tool
+
+  def test_get_declaration(self, mock_rest_api_tool):
+    """Test _get_declaration method."""
+    tool = GoogleApiTool(mock_rest_api_tool)
+
+    declaration = tool._get_declaration()
+
+    assert isinstance(declaration, FunctionDeclaration)
+    assert declaration.name == "test_function"
+    assert declaration.description == "Test function description"
+    mock_rest_api_tool._get_declaration.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_run_async(self, mock_rest_api_tool, mock_tool_context):
+    """Test run_async method."""
+    tool = GoogleApiTool(mock_rest_api_tool)
+    args = {"param1": "value1"}
+
+    result = await tool.run_async(args=args, tool_context=mock_tool_context)
+
+    assert result == {"result": "success"}
+    mock_rest_api_tool.run_async.assert_called_once_with(
+        args=args, tool_context=mock_tool_context
+    )
+
+  def test_configure_auth(self, mock_rest_api_tool):
+    """Test configure_auth method."""
+    tool = GoogleApiTool(mock_rest_api_tool)
+    client_id = "test_client_id"
+    client_secret = "test_client_secret"
+
+    tool.configure_auth(client_id=client_id, client_secret=client_secret)
+
+    # Check that auth_credential was set correctly on the rest_api_tool
+    assert mock_rest_api_tool.auth_credential is not None
+    assert (
+        mock_rest_api_tool.auth_credential.auth_type
+        == AuthCredentialTypes.OPEN_ID_CONNECT
+    )
+    assert mock_rest_api_tool.auth_credential.oauth2.client_id == client_id
+    assert (
+        mock_rest_api_tool.auth_credential.oauth2.client_secret == client_secret
+    )
+
+  @mock.patch(
+      "google.adk.tools.google_api_tool.google_api_tool.service_account_scheme_credential"
+  )
+  def test_configure_sa_auth(
+      self, mock_service_account_scheme_credential, mock_rest_api_tool
+  ):
+    """Test configure_sa_auth method."""
+    # Setup mock return values
+    mock_auth_scheme = mock.MagicMock()
+    mock_auth_credential = mock.MagicMock()
+    mock_service_account_scheme_credential.return_value = (
+        mock_auth_scheme,
+        mock_auth_credential,
+    )
+
+    service_account = ServiceAccount(
+        service_account_credential=ServiceAccountCredential(
+            type="service_account",
+            project_id="project_id",
+            private_key_id="private_key_id",
+            private_key="private_key",
+            client_email="client_email",
+            client_id="client_id",
+            auth_uri="auth_uri",
+            token_uri="token_uri",
+            auth_provider_x509_cert_url="auth_provider_x509_cert_url",
+            client_x509_cert_url="client_x509_cert_url",
+            universe_domain="universe_domain",
+        ),
+        scopes=["scope1", "scope2"],
+    )
+
+    # Create tool and call method
+    tool = GoogleApiTool(mock_rest_api_tool)
+    tool.configure_sa_auth(service_account=service_account)
+
+    # Verify service_account_scheme_credential was called correctly
+    mock_service_account_scheme_credential.assert_called_once_with(
+        service_account
+    )
+
+    # Verify auth_scheme and auth_credential were set correctly on the rest_api_tool
+    assert mock_rest_api_tool.auth_scheme == mock_auth_scheme
+    assert mock_rest_api_tool.auth_credential == mock_auth_credential