chore: import upstream snapshot with attribution

tests/unittests/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/utils/test_cache_performance_analyzer.py

@@ -0,0 +1,494 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for CachePerformanceAnalyzer."""
+
+import time
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+
+from google.adk.events.event import Event
+from google.adk.models.cache_metadata import CacheMetadata
+from google.adk.sessions.base_session_service import BaseSessionService
+from google.adk.sessions.session import Session
+from google.adk.utils.cache_performance_analyzer import CachePerformanceAnalyzer
+from google.genai import types
+import pytest
+
+
+class TestCachePerformanceAnalyzer:
+  """Test suite for CachePerformanceAnalyzer."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.mock_session_service = MagicMock(spec=BaseSessionService)
+    self.analyzer = CachePerformanceAnalyzer(self.mock_session_service)
+
+  def create_cache_metadata(
+      self, invocations_used=1, cache_name="test-cache", contents_count=5
+  ):
+    """Helper to create test CacheMetadata."""
+    return CacheMetadata(
+        cache_name=(
+            f"projects/test/locations/us-central1/cachedContents/{cache_name}"
+        ),
+        expire_time=time.time() + 1800,
+        fingerprint="test_fingerprint",
+        invocations_used=invocations_used,
+        contents_count=contents_count,
+        created_at=time.time() - 600,
+    )
+
+  def create_mock_usage_metadata(
+      self, prompt_tokens=1000, cached_tokens=500, candidates_tokens=100
+  ):
+    """Helper to create mock usage metadata."""
+    return types.GenerateContentResponseUsageMetadata(
+        prompt_token_count=prompt_tokens,
+        cached_content_token_count=cached_tokens,
+        candidates_token_count=candidates_tokens,
+        total_token_count=prompt_tokens + candidates_tokens,
+    )
+
+  def create_mock_event(
+      self, author="test_agent", cache_metadata=None, usage_metadata=None
+  ):
+    """Helper to create mock event."""
+    event = Event(author=author, cache_metadata=cache_metadata)
+    if usage_metadata:
+      event.usage_metadata = usage_metadata
+    return event
+
+  def test_init(self):
+    """Test analyzer initialization."""
+    assert self.analyzer.session_service == self.mock_session_service
+
+  async def test_get_agent_cache_history_empty_session(self):
+    """Test getting cache history from empty session."""
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=[],
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer._get_agent_cache_history(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    assert result == []
+
+  async def test_get_agent_cache_history_no_cache_events(self):
+    """Test getting cache history when no events have cache metadata."""
+    events = [
+        self.create_mock_event(author="test_agent"),
+        self.create_mock_event(author="other_agent"),
+        self.create_mock_event(author="test_agent"),
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer._get_agent_cache_history(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    assert result == []
+
+  async def test_get_agent_cache_history_specific_agent(self):
+    """Test getting cache history for specific agent."""
+    cache1 = self.create_cache_metadata(invocations_used=1, cache_name="cache1")
+    cache2 = self.create_cache_metadata(invocations_used=3, cache_name="cache2")
+    cache3 = self.create_cache_metadata(invocations_used=5, cache_name="cache3")
+
+    events = [
+        self.create_mock_event(author="test_agent", cache_metadata=cache1),
+        self.create_mock_event(author="other_agent", cache_metadata=cache2),
+        self.create_mock_event(author="test_agent", cache_metadata=cache3),
+        self.create_mock_event(author="test_agent"),  # No cache metadata
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer._get_agent_cache_history(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    # Should only return cache metadata for test_agent
+    assert len(result) == 2
+    assert result[0] == cache1
+    assert result[1] == cache3
+
+  async def test_get_agent_cache_history_all_agents(self):
+    """Test getting cache history for all agents."""
+    cache1 = self.create_cache_metadata(invocations_used=1, cache_name="cache1")
+    cache2 = self.create_cache_metadata(invocations_used=3, cache_name="cache2")
+
+    events = [
+        self.create_mock_event(author="agent1", cache_metadata=cache1),
+        self.create_mock_event(author="agent2", cache_metadata=cache2),
+        self.create_mock_event(author="agent1"),  # No cache metadata
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    # Pass None for agent_name to get all agents
+    result = await self.analyzer._get_agent_cache_history(
+        "test_session", "test_user", "test_app", None
+    )
+
+    # Should return cache metadata for all agents
+    assert len(result) == 2
+    assert result[0] == cache1
+    assert result[1] == cache2
+
+  async def test_analyze_agent_cache_performance_no_cache_data(self):
+    """Test analysis with no cache data."""
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=[],
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    assert result["status"] == "no_cache_data"
+
+  async def test_analyze_agent_cache_performance_with_cache_data(self):
+    """Test comprehensive analysis with cache data and token metrics."""
+    cache1 = self.create_cache_metadata(invocations_used=2, cache_name="cache1")
+    cache2 = self.create_cache_metadata(invocations_used=5, cache_name="cache2")
+    cache3 = self.create_cache_metadata(invocations_used=8, cache_name="cache3")
+
+    usage1 = self.create_mock_usage_metadata(
+        prompt_tokens=1000, cached_tokens=800
+    )
+    usage2 = self.create_mock_usage_metadata(
+        prompt_tokens=1500, cached_tokens=1200
+    )
+    usage3 = self.create_mock_usage_metadata(prompt_tokens=800, cached_tokens=0)
+
+    events = [
+        self.create_mock_event(
+            author="test_agent", cache_metadata=cache1, usage_metadata=usage1
+        ),
+        self.create_mock_event(author="other_agent", cache_metadata=cache2),
+        self.create_mock_event(
+            author="test_agent", cache_metadata=cache2, usage_metadata=usage2
+        ),
+        self.create_mock_event(
+            author="test_agent", cache_metadata=cache3, usage_metadata=usage3
+        ),
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    # Basic cache metrics
+    assert result["status"] == "active"
+    assert result["requests_with_cache"] == 3
+    assert result["cache_refreshes"] == 3  # 3 unique cache names
+    assert result["total_invocations"] == 15  # 2 + 5 + 8
+
+    expected_avg_invocations = (2 + 5 + 8) / 3  # 5.0
+    assert result["avg_invocations_used"] == expected_avg_invocations
+
+    # Token metrics
+    assert result["total_prompt_tokens"] == 3300  # 1000 + 1500 + 800
+    assert result["total_cached_tokens"] == 2000  # 800 + 1200 + 0
+    assert result["total_requests"] == 3
+    assert (
+        result["requests_with_cache_hits"] == 2
+    )  # Only first two have cached tokens
+
+    # Calculated metrics
+    expected_hit_ratio = (2000 / 3300) * 100  # ~60.6%
+    expected_utilization = (2 / 3) * 100  # ~66.7%
+    expected_avg_cached = 2000 / 3  # ~666.7
+
+    assert abs(result["cache_hit_ratio_percent"] - expected_hit_ratio) < 0.01
+    assert (
+        abs(result["cache_utilization_ratio_percent"] - expected_utilization)
+        < 0.01
+    )
+    assert (
+        abs(result["avg_cached_tokens_per_request"] - expected_avg_cached)
+        < 0.01
+    )
+
+  async def test_analyze_agent_cache_performance_single_cache(self):
+    """Test analysis with single cache instance."""
+    cache = self.create_cache_metadata(
+        invocations_used=10, cache_name="single_cache"
+    )
+    usage = self.create_mock_usage_metadata(
+        prompt_tokens=2000, cached_tokens=1500
+    )
+
+    events = [
+        self.create_mock_event(
+            author="test_agent", cache_metadata=cache, usage_metadata=usage
+        ),
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    assert result["status"] == "active"
+    assert result["requests_with_cache"] == 1
+    assert result["avg_invocations_used"] == 10.0
+    assert result["cache_refreshes"] == 1
+    assert result["total_invocations"] == 10
+    assert result["latest_cache"] == cache.cache_name
+
+    # Token metrics for single request
+    assert result["total_prompt_tokens"] == 2000
+    assert result["total_cached_tokens"] == 1500
+    assert result["cache_hit_ratio_percent"] == 75.0  # 1500/2000 * 100
+    assert result["cache_utilization_ratio_percent"] == 100.0  # 1/1 * 100
+    assert result["avg_cached_tokens_per_request"] == 1500.0
+
+  async def test_analyze_agent_cache_performance_no_token_data(self):
+    """Test analysis when events have no usage_metadata."""
+    cache = self.create_cache_metadata(invocations_used=5)
+
+    events = [
+        self.create_mock_event(author="test_agent", cache_metadata=cache),
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    # Should still work but with zero token metrics
+    assert result["status"] == "active"
+    assert result["requests_with_cache"] == 1
+    assert result["total_prompt_tokens"] == 0
+    assert result["total_cached_tokens"] == 0
+    assert result["cache_hit_ratio_percent"] == 0.0
+    assert result["cache_utilization_ratio_percent"] == 0.0
+    assert result["avg_cached_tokens_per_request"] == 0.0
+
+  async def test_analyze_agent_cache_performance_zero_invocations(self):
+    """Test analysis with zero invocations."""
+    cache = self.create_cache_metadata(
+        invocations_used=0, cache_name="zero_cache"
+    )
+    usage = self.create_mock_usage_metadata(
+        prompt_tokens=1000, cached_tokens=500
+    )
+
+    events = [
+        self.create_mock_event(
+            author="test_agent", cache_metadata=cache, usage_metadata=usage
+        ),
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    assert result["status"] == "active"
+    assert result["avg_invocations_used"] == 0.0
+    assert result["total_invocations"] == 0
+
+    # Token metrics should still work
+    assert result["total_prompt_tokens"] == 1000
+    assert result["total_cached_tokens"] == 500
+
+  async def test_session_service_integration(self):
+    """Test integration with session service."""
+    cache_metadata = self.create_cache_metadata(invocations_used=7)
+
+    events = [
+        self.create_mock_event(
+            author="integration_agent", cache_metadata=cache_metadata
+        ),
+    ]
+
+    mock_session = Session(
+        id="integration_session",
+        app_name="integration_app",
+        user_id="integration_user",
+        events=events,
+    )
+
+    # Configure the mock to return the session
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "integration_session",
+        "integration_user",
+        "integration_app",
+        "integration_agent",
+    )
+
+    # Verify the session service was called with correct parameters (twice internally)
+    assert self.mock_session_service.get_session.call_count == 2
+    self.mock_session_service.get_session.assert_called_with(
+        session_id="integration_session",
+        app_name="integration_app",
+        user_id="integration_user",
+    )
+
+    assert result["status"] == "active"
+    assert result["requests_with_cache"] == 1
+
+  async def test_analyze_agent_cache_performance_with_fingerprint_only(self):
+    """Fingerprint-only entries (cache_name=None, invocations_used=None) don't crash."""
+    fp_only = CacheMetadata(fingerprint="fp", contents_count=3)
+    active = self.create_cache_metadata(invocations_used=4, cache_name="active")
+    fp_usage = self.create_mock_usage_metadata(
+        prompt_tokens=1000, cached_tokens=0
+    )
+    active_usage = self.create_mock_usage_metadata(
+        prompt_tokens=1000, cached_tokens=800
+    )
+
+    events = [
+        self.create_mock_event(
+            author="test_agent",
+            cache_metadata=fp_only,
+            usage_metadata=fp_usage,
+        ),
+        self.create_mock_event(
+            author="test_agent",
+            cache_metadata=active,
+            usage_metadata=active_usage,
+        ),
+    ]
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "test_agent"
+    )
+
+    assert result["status"] == "active"
+    assert result["total_requests"] == 2
+    assert result["total_prompt_tokens"] == 2000
+    assert result["total_cached_tokens"] == 800
+    assert result["total_invocations"] == 4
+    assert result["avg_invocations_used"] == 4.0
+    assert result["cache_refreshes"] == 1
+    assert result["requests_with_cache"] == 2
+
+  async def test_mixed_agents_filtering(self):
+    """Test that analysis correctly filters by agent name."""
+    target_cache = self.create_cache_metadata(
+        invocations_used=3, cache_name="target"
+    )
+    other_cache = self.create_cache_metadata(
+        invocations_used=5, cache_name="other"
+    )
+
+    target_usage = self.create_mock_usage_metadata(
+        prompt_tokens=1000, cached_tokens=800
+    )
+    other_usage = self.create_mock_usage_metadata(
+        prompt_tokens=2000, cached_tokens=1600
+    )
+
+    events = [
+        self.create_mock_event(
+            author="target_agent",
+            cache_metadata=target_cache,
+            usage_metadata=target_usage,
+        ),
+        self.create_mock_event(
+            author="other_agent",
+            cache_metadata=other_cache,
+            usage_metadata=other_usage,
+        ),
+        self.create_mock_event(author="target_agent"),  # No cache data
+    ]
+
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=events,
+    )
+    self.mock_session_service.get_session = AsyncMock(return_value=mock_session)
+
+    result = await self.analyzer.analyze_agent_cache_performance(
+        "test_session", "test_user", "test_app", "target_agent"
+    )
+
+    # Should only include target_agent's data
+    assert result["requests_with_cache"] == 1
+    assert result["total_invocations"] == 3
+    assert result["total_prompt_tokens"] == 1000  # Only target_agent's tokens
+    assert result["total_cached_tokens"] == 800  # Only target_agent's tokens

tests/unittests/utils/test_client_labels_utils.py

@@ -0,0 +1,68 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+from google.adk import version
+from google.adk.utils import _client_labels_utils
+import pytest
+
+
+def test_get_client_labels_default():
+  """Test get_client_labels returns default labels."""
+  labels = _client_labels_utils.get_client_labels()
+  assert len(labels) == 2
+  assert f"google-adk/{version.__version__}" == labels[0]
+  assert f"gl-python/{sys.version.split()[0]}" == labels[1]
+
+
+def test_get_client_labels_with_agent_engine_id(monkeypatch):
+  """Test get_client_labels returns agent engine tag when env var is set."""
+  monkeypatch.setenv(
+      _client_labels_utils._AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME,
+      "test-agent-id",
+  )
+  labels = _client_labels_utils.get_client_labels()
+  assert len(labels) == 2
+  assert (
+      f"google-adk/{version.__version__}+{_client_labels_utils._AGENT_ENGINE_TELEMETRY_TAG}"
+      == labels[0]
+  )
+  assert f"gl-python/{sys.version.split()[0]}" == labels[1]
+
+
+def test_get_client_labels_with_context():
+  """Test get_client_labels includes label from context."""
+  with _client_labels_utils.client_label_context("my-label/1.0"):
+    labels = _client_labels_utils.get_client_labels()
+    assert len(labels) == 3
+    assert f"google-adk/{version.__version__}" == labels[0]
+    assert f"gl-python/{sys.version.split()[0]}" == labels[1]
+    assert "my-label/1.0" == labels[2]
+
+
+def test_client_label_context_nested_error():
+  """Test client_label_context raises error when nested."""
+  with pytest.raises(ValueError, match="Client label already exists"):
+    with _client_labels_utils.client_label_context("my-label/1.0"):
+      with _client_labels_utils.client_label_context("another-label/1.0"):
+        pass
+
+
+def test_eval_client_label():
+  """Test EVAL_CLIENT_LABEL has correct format."""
+  assert (
+      f"google-adk-eval/{version.__version__}"
+      == _client_labels_utils.EVAL_CLIENT_LABEL
+  )

tests/unittests/utils/test_content_utils.py

@@ -0,0 +1,68 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.utils.content_utils import SKIP_THOUGHT_SIGNATURE_VALIDATOR
+from google.adk.utils.content_utils import to_user_content
+from google.genai import types
+from pydantic import BaseModel
+
+
+def test_skip_thought_signature_validator_wire_value():
+  # The backend recognizes this exact byte string to bypass validation;
+  # changing it would break every replayed synthetic part.
+  assert SKIP_THOUGHT_SIGNATURE_VALIDATOR == b'skip_thought_signature_validator'
+
+
+def test_skip_thought_signature_validator_assignable_to_part():
+  part = types.Part(
+      text='injected',
+      thought_signature=SKIP_THOUGHT_SIGNATURE_VALIDATOR,
+  )
+  assert part.thought_signature == SKIP_THOUGHT_SIGNATURE_VALIDATOR
+
+
+def test_to_user_content_str_input_becomes_user_text():
+  content = to_user_content('hello')
+  assert content.role == 'user'
+  assert content.parts[0].text == 'hello'
+
+
+def test_to_user_content_input_is_normalized_to_user_role():
+  original = types.Content(role='model', parts=[types.Part(text='hi')])
+  content = to_user_content(original)
+  assert content.role == 'user'
+  assert content.parts[0].text == 'hi'
+
+
+def test_to_user_content_basemodel_input_is_json():
+  class _M(BaseModel):
+    a: int
+
+  content = to_user_content(_M(a=1))
+  assert content.role == 'user'
+  assert '"a":1' in content.parts[0].text.replace(' ', '')
+
+
+def test_to_user_content_dict_input_is_json():
+  content = to_user_content({'a': 1})
+  assert content.role == 'user'
+  assert content.parts[0].text.replace(' ', '') == '{"a":1}'
+
+
+def test_to_user_content_other_input_is_str():
+  content = to_user_content(42)
+  assert content.role == 'user'
+  assert content.parts[0].text == '42'

tests/unittests/utils/test_context_utils.py

@@ -0,0 +1,155 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for context_utils module."""
+
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.context import Context
+from google.adk.tools.tool_context import ToolContext
+from google.adk.utils import context_utils
+from google.adk.utils.context_utils import find_context_parameter
+
+
+class TestFindContextParameter:
+  """Tests for find_context_parameter function."""
+
+  def test_find_context_parameter_with_context_type(self):
+    """Test detection of Context type annotation."""
+
+    def my_tool(query: str, ctx: Context) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'ctx'
+
+  def test_find_context_parameter_with_string_annotation(self):
+    """Test detection of string annotation 'Context'."""
+
+    def my_tool(query: str, ctx: 'Context') -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'ctx'
+
+  def test_find_context_parameter_with_string_tool_context(self):
+    """Test detection of string annotation 'ToolContext'."""
+
+    def my_tool(query: str, ctx: 'ToolContext') -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'ctx'
+
+  def test_find_context_parameter_with_string_optional_context(self):
+    """Test detection of string annotation 'Optional[Context]'."""
+
+    def my_tool(query: str, ctx: 'Optional[Context]' = None) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'ctx'
+
+  def test_find_context_parameter_with_tool_context_type(self):
+    """Test detection of ToolContext type annotation."""
+
+    def my_tool(query: str, tool_context: ToolContext) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'tool_context'
+
+  def test_find_context_parameter_with_callback_context_type(self):
+    """Test detection of CallbackContext type annotation."""
+
+    def my_callback(ctx: CallbackContext) -> None:
+      pass
+
+    assert find_context_parameter(my_callback) == 'ctx'
+
+  def test_find_context_parameter_with_optional_context(self):
+    """Test detection of Optional[Context] type annotation."""
+
+    def my_tool(query: str, context: Optional[Context] = None) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'context'
+
+  def test_find_context_parameter_with_custom_name(self):
+    """Test that any parameter name works with Context type."""
+
+    def my_tool(query: str, my_custom_ctx: Context) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'my_custom_ctx'
+
+  def test_find_context_parameter_no_context(self):
+    """Test function without context parameter returns None."""
+
+    def my_tool(query: str, count: int) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) is None
+
+  def test_find_context_parameter_no_annotations(self):
+    """Test function without type annotations returns None."""
+
+    def my_tool(query, ctx):
+      return query
+
+    assert find_context_parameter(my_tool) is None
+
+  def test_find_context_parameter_with_none_func(self):
+    """Test that None function returns None."""
+    assert find_context_parameter(None) is None
+
+  def test_find_context_parameter_returns_first_match(self):
+    """Test that first context parameter is returned if multiple exist."""
+
+    def my_tool(first_ctx: Context, second_ctx: Context) -> str:
+      return 'test'
+
+    assert find_context_parameter(my_tool) == 'first_ctx'
+
+  def test_find_context_parameter_with_mixed_params(self):
+    """Test context parameter detection with various other parameters."""
+
+    def my_tool(
+        query: str,
+        count: int,
+        ctx: Context,
+        optional_param: Optional[str] = None,
+    ) -> str:
+      return query
+
+    assert find_context_parameter(my_tool) == 'ctx'
+
+
+class TestFindContextParameterCaching:
+  """Tests for find_context_parameter caching behavior."""
+
+  def test_repeated_calls_inspect_signature_once(self):
+    """Repeated calls with the same function reuse the cached result."""
+
+    def my_tool(ctx: Context) -> str:
+      return 'ok'
+
+    find_context_parameter.cache_clear()
+
+    with mock.patch.object(
+        context_utils.inspect,
+        'signature',
+        wraps=context_utils.inspect.signature,
+    ) as spy:
+      for _ in range(10):
+        assert find_context_parameter(my_tool) == 'ctx'
+
+    assert spy.call_count == 1

tests/unittests/utils/test_env_utils.py

@@ -0,0 +1,84 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+
+from google.adk.utils.env_utils import is_enterprise_mode_enabled
+from google.adk.utils.env_utils import is_env_enabled
+import pytest
+
+
+@pytest.mark.parametrize(
+    'env_value,expected',
+    [
+        ('true', True),
+        ('TRUE', True),
+        ('TrUe', True),
+        ('1', True),
+        ('false', False),
+        ('FALSE', False),
+        ('0', False),
+        ('', False),
+    ],
+)
+def test_is_env_enabled(monkeypatch, env_value, expected):
+  """Test is_env_enabled with various environment variable values."""
+  monkeypatch.setenv('TEST_FLAG', env_value)
+  assert is_env_enabled('TEST_FLAG') is expected
+
+
+@pytest.mark.parametrize(
+    'default,expected',
+    [
+        ('0', False),
+        ('1', True),
+        ('true', True),
+    ],
+)
+def test_is_env_enabled_with_defaults(monkeypatch, default, expected):
+  """Test is_env_enabled when env var is not set with different defaults."""
+  monkeypatch.delenv('TEST_FLAG', raising=False)
+  assert is_env_enabled('TEST_FLAG', default=default) is expected
+
+
+def test_is_enterprise_mode_enabled_via_enterprise_env(monkeypatch):
+  """Enterprise mode is on when GOOGLE_GENAI_USE_ENTERPRISE is truthy."""
+  monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true')
+
+  assert is_enterprise_mode_enabled() is True
+
+
+def test_is_enterprise_mode_enabled_falls_back_to_vertexai_with_warning(
+    monkeypatch,
+):
+  """The deprecated GOOGLE_GENAI_USE_VERTEXAI still enables enterprise mode and warns."""
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+
+  with warnings.catch_warnings(record=True) as caught:
+    warnings.simplefilter('always')
+    result = is_enterprise_mode_enabled()
+
+  assert result is True
+  assert len(caught) == 1
+  assert issubclass(caught[-1].category, DeprecationWarning)
+  assert 'GOOGLE_GENAI_USE_VERTEXAI is deprecated' in str(caught[-1].message)
+
+
+def test_is_enterprise_mode_enabled_defaults_to_false(monkeypatch):
+  """Enterprise mode is off when no relevant env var is set."""
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.delenv('GOOGLE_GENAI_USE_VERTEXAI', raising=False)
+
+  assert is_enterprise_mode_enabled() is False

tests/unittests/utils/test_feature_decorator.py

@@ -0,0 +1,384 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+import warnings
+
+from google.adk.utils.feature_decorator import experimental
+from google.adk.utils.feature_decorator import working_in_progress
+
+
+@working_in_progress("in complete feature, don't use yet")
+class IncompleteFeature:
+
+  def run(self):
+    return "running"
+
+
+@working_in_progress("function not ready")
+def wip_function():
+  return "executing"
+
+
+@experimental("api may have breaking change in the future.")
+def experimental_fn():
+  return "executing"
+
+
+@experimental("class may change")
+class ExperimentalClass:
+
+  def run(self):
+    return "running experimental"
+
+
+# Test classes/functions for new usage patterns
+@experimental
+class ExperimentalClassNoParens:
+
+  def run(self):
+    return "running experimental without parens"
+
+
+@experimental()
+class ExperimentalClassEmptyParens:
+
+  def run(self):
+    return "running experimental with empty parens"
+
+
+@experimental
+def experimental_fn_no_parens():
+  return "executing without parens"
+
+
+@experimental()
+def experimental_fn_empty_parens():
+  return "executing with empty parens"
+
+
+def test_working_in_progress_class_raises_error():
+  """Test that WIP class raises RuntimeError by default."""
+  # Ensure environment variable is not set
+  if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+    del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+  try:
+    feature = IncompleteFeature()
+    assert False, "Expected RuntimeError to be raised"
+  except RuntimeError as e:
+    assert "[WIP] IncompleteFeature:" in str(e)
+    assert "don't use yet" in str(e)
+
+
+def test_working_in_progress_function_raises_error():
+  """Test that WIP function raises RuntimeError by default."""
+  # Ensure environment variable is not set
+  if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+    del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+  try:
+    result = wip_function()
+    assert False, "Expected RuntimeError to be raised"
+  except RuntimeError as e:
+    assert "[WIP] wip_function:" in str(e)
+    assert "function not ready" in str(e)
+
+
+def test_working_in_progress_class_bypassed_with_env_var():
+  """Test that WIP class works without warnings when env var is set."""
+  # Set the bypass environment variable
+  os.environ["ADK_ALLOW_WIP_FEATURES"] = "true"
+
+  try:
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+
+      feature = IncompleteFeature()
+      result = feature.run()
+
+      assert result == "running"
+      # Should have no warnings when bypassed
+      assert len(w) == 0
+  finally:
+    # Clean up environment variable
+    if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+      del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+
+def test_working_in_progress_function_bypassed_with_env_var():
+  """Test that WIP function works without warnings when env var is set."""
+  # Set the bypass environment variable
+  os.environ["ADK_ALLOW_WIP_FEATURES"] = "true"
+
+  try:
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+
+      result = wip_function()
+
+      assert result == "executing"
+      # Should have no warnings when bypassed
+      assert len(w) == 0
+  finally:
+    # Clean up environment variable
+    if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+      del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+
+def test_working_in_progress_env_var_case_insensitive():
+  """Test that WIP bypass works with different case values."""
+  test_cases = ["true", "True", "TRUE", "tRuE"]
+
+  for case in test_cases:
+    os.environ["ADK_ALLOW_WIP_FEATURES"] = case
+
+    try:
+      with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+
+        result = wip_function()
+
+        assert result == "executing"
+        assert len(w) == 0
+    finally:
+      if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+        del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+
+def test_working_in_progress_env_var_false_values():
+  """Test that WIP still raises errors with false-like env var values."""
+  false_values = ["false", "False", "FALSE", "0", "", "anything_else"]
+
+  for false_val in false_values:
+    os.environ["ADK_ALLOW_WIP_FEATURES"] = false_val
+
+    try:
+      result = wip_function()
+      assert False, f"Expected RuntimeError with env var '{false_val}'"
+    except RuntimeError as e:
+      assert "[WIP] wip_function:" in str(e)
+    finally:
+      if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+        del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+
+def test_working_in_progress_loads_from_dotenv_file():
+  """Test that WIP decorator can load environment variables from .env file."""
+  # Skip test if dotenv is not available
+  try:
+    from dotenv import load_dotenv
+  except ImportError:
+    import pytest
+
+    pytest.skip("python-dotenv not available")
+
+  # Ensure environment variable is not set in os.environ
+  if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+    del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+  # Create a temporary .env file in current directory
+  dotenv_path = ".env.test"
+
+  try:
+    # Write the env file
+    with open(dotenv_path, "w") as f:
+      f.write("ADK_ALLOW_WIP_FEATURES=true\n")
+
+    # Load the environment variables from the file
+    load_dotenv(dotenv_path)
+
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+
+      # This should work because the .env file contains ADK_ALLOW_WIP_FEATURES=true
+      result = wip_function()
+
+      assert result == "executing"
+      # Should have no warnings when bypassed via .env file
+      assert len(w) == 0
+
+  finally:
+    # Clean up
+    try:
+      os.unlink(dotenv_path)
+    except FileNotFoundError:
+      pass
+    if "ADK_ALLOW_WIP_FEATURES" in os.environ:
+      del os.environ["ADK_ALLOW_WIP_FEATURES"]
+
+
+def test_experimental_function_warns(monkeypatch):
+  """Test that experimental function shows warnings (unchanged behavior)."""
+  # Ensure environment variable is not set
+  monkeypatch.delenv(
+      "ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
+  )
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+
+    result = experimental_fn()
+
+    assert result == "executing"
+    assert len(w) == 1
+    assert issubclass(w[0].category, UserWarning)
+    assert "[EXPERIMENTAL] experimental_fn:" in str(w[0].message)
+    assert "breaking change in the future" in str(w[0].message)
+
+
+def test_experimental_class_warns(monkeypatch):
+  """Test that experimental class shows warnings (unchanged behavior)."""
+  # Ensure environment variable is not set
+  monkeypatch.delenv(
+      "ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
+  )
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+
+    exp_class = ExperimentalClass()
+    result = exp_class.run()
+
+    assert result == "running experimental"
+    assert len(w) == 1
+    assert issubclass(w[0].category, UserWarning)
+    assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message)
+    assert "class may change" in str(w[0].message)
+
+
+def test_experimental_function_bypassed_with_env_var(monkeypatch):
+  """Experimental function emits no warning when bypass env var is true."""
+  true_values = ["true", "True", "TRUE", "1", "yes", "YES", "on", "ON"]
+  for true_val in true_values:
+    monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", true_val)
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+      result = experimental_fn()
+      assert result == "executing"
+      assert len(w) == 0, f"Bypass failed for env value {true_val}"
+
+
+def test_experimental_class_bypassed_with_env_var(monkeypatch):
+  """Experimental class emits no warning when bypass env var is true."""
+  true_values = ["true", "True", "TRUE", "1", "yes", "YES", "on", "ON"]
+  for true_val in true_values:
+    monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", true_val)
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+      exp_class = ExperimentalClass()
+      result = exp_class.run()
+      assert result == "running experimental"
+      assert len(w) == 0, f"Bypass failed for env value {true_val}"
+
+
+def test_experimental_function_not_bypassed_for_false_env_var(monkeypatch):
+  """Experimental function still warns for non-true bypass env var values."""
+  false_values = ["false", "False", "FALSE", "0", "", "no", "off"]
+  for false_val in false_values:
+    monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", false_val)
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+      experimental_fn()
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] experimental_fn:" in str(w[0].message)
+
+
+def test_experimental_class_not_bypassed_for_false_env_var(monkeypatch):
+  """Experimental class still warns for non-true bypass env var values."""
+  false_values = ["false", "False", "FALSE", "0", "", "no", "off"]
+  for false_val in false_values:
+    monkeypatch.setenv("ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", false_val)
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter("always")
+      ExperimentalClass()
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message)
+
+
+def test_experimental_class_no_parens_warns(monkeypatch):
+  """Test that experimental class without parentheses shows default warning."""
+  monkeypatch.delenv(
+      "ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
+  )
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+
+    exp_class = ExperimentalClassNoParens()
+    result = exp_class.run()
+
+    assert result == "running experimental without parens"
+    assert len(w) == 1
+    assert issubclass(w[0].category, UserWarning)
+    assert "[EXPERIMENTAL] ExperimentalClassNoParens:" in str(w[0].message)
+    assert "This feature is experimental and may change or be removed" in str(
+        w[0].message
+    )
+
+
+def test_experimental_class_empty_parens_warns(monkeypatch):
+  """Test that experimental class with empty parentheses shows default warning."""
+  monkeypatch.delenv(
+      "ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
+  )
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+
+    exp_class = ExperimentalClassEmptyParens()
+    result = exp_class.run()
+
+    assert result == "running experimental with empty parens"
+    assert len(w) == 1
+    assert issubclass(w[0].category, UserWarning)
+    assert "[EXPERIMENTAL] ExperimentalClassEmptyParens:" in str(w[0].message)
+    assert "This feature is experimental and may change or be removed" in str(
+        w[0].message
+    )
+
+
+def test_experimental_function_no_parens_warns(monkeypatch):
+  """Test that experimental function without parentheses shows default warning."""
+  monkeypatch.delenv(
+      "ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
+  )
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+
+    result = experimental_fn_no_parens()
+
+    assert result == "executing without parens"
+    assert len(w) == 1
+    assert issubclass(w[0].category, UserWarning)
+    assert "[EXPERIMENTAL] experimental_fn_no_parens:" in str(w[0].message)
+    assert "This feature is experimental and may change or be removed" in str(
+        w[0].message
+    )
+
+
+def test_experimental_function_empty_parens_warns(monkeypatch):
+  """Test that experimental function with empty parentheses shows default warning."""
+  monkeypatch.delenv(
+      "ADK_SUPPRESS_EXPERIMENTAL_FEATURE_WARNINGS", raising=False
+  )
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+
+    result = experimental_fn_empty_parens()
+
+    assert result == "executing with empty parens"
+    assert len(w) == 1
+    assert issubclass(w[0].category, UserWarning)
+    assert "[EXPERIMENTAL] experimental_fn_empty_parens:" in str(w[0].message)
+    assert "This feature is experimental and may change or be removed" in str(
+        w[0].message
+    )

tests/unittests/utils/test_google_client_headers.py

@@ -0,0 +1,88 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+from google.adk import version
+from google.adk.utils import _google_client_headers
+import pytest
+
+_EXPECTED_BASE_HEADER = (
+    f"google-adk/{version.__version__} gl-python/{sys.version.split()[0]}"
+)
+
+
+def test_get_tracking_headers():
+  """Test get_tracking_headers returns correct headers."""
+  headers = _google_client_headers.get_tracking_headers()
+  assert headers == {
+      "x-goog-api-client": _EXPECTED_BASE_HEADER,
+      "user-agent": _EXPECTED_BASE_HEADER,
+  }
+
+
+@pytest.mark.parametrize(
+    "input_headers, expected_headers",
+    [
+        (
+            None,
+            {
+                "x-goog-api-client": _EXPECTED_BASE_HEADER,
+                "user-agent": _EXPECTED_BASE_HEADER,
+            },
+        ),
+        (
+            {},
+            {
+                "x-goog-api-client": _EXPECTED_BASE_HEADER,
+                "user-agent": _EXPECTED_BASE_HEADER,
+            },
+        ),
+        (
+            {"x-goog-api-client": "label3 label4"},
+            {
+                "x-goog-api-client": f"{_EXPECTED_BASE_HEADER} label3 label4",
+                "user-agent": _EXPECTED_BASE_HEADER,
+            },
+        ),
+        (
+            {"x-goog-api-client": f"gl-python/{sys.version.split()[0]} label3"},
+            {
+                "x-goog-api-client": f"{_EXPECTED_BASE_HEADER} label3",
+                "user-agent": _EXPECTED_BASE_HEADER,
+            },
+        ),
+        (
+            {"other-header": "value"},
+            {
+                "x-goog-api-client": _EXPECTED_BASE_HEADER,
+                "user-agent": _EXPECTED_BASE_HEADER,
+                "other-header": "value",
+            },
+        ),
+    ],
+)
+def test_merge_tracking_headers(input_headers, expected_headers):
+  """Test merge_tracking_headers with various inputs."""
+  headers = _google_client_headers.merge_tracking_headers(input_headers)
+  assert headers == expected_headers
+
+
+def test_get_tracking_http_options():
+  """get_tracking_http_options returns HttpOptions carrying tracking headers."""
+  http_options = _google_client_headers.get_tracking_http_options()
+  assert http_options.headers == {
+      "x-goog-api-client": _EXPECTED_BASE_HEADER,
+      "user-agent": _EXPECTED_BASE_HEADER,
+  }

tests/unittests/utils/test_instructions_utils.py

@@ -0,0 +1,282 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.sessions.session import Session
+from google.adk.utils import instructions_utils
+import pytest
+
+from .. import testing_utils
+
+
+class MockArtifactService:
+
+  def __init__(self, artifacts: dict):
+    self.artifacts = artifacts
+
+  async def load_artifact(self, app_name, user_id, session_id, filename):
+    if filename in self.artifacts:
+      return self.artifacts[filename]
+    else:
+      return None
+
+
+async def _create_test_readonly_context(
+    state: dict = None,
+    artifact_service: MockArtifactService = None,
+    app_name: str = "test_app",
+    user_id: str = "test_user",
+    session_id: str = "test_session_id",
+) -> ReadonlyContext:
+  agent = Agent(
+      model="gemini-2.5-flash",
+      name="agent",
+      instruction="test",
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  invocation_context.session = Session(
+      state=state if state else {},
+      app_name=app_name,
+      user_id=user_id,
+      id=session_id,
+  )
+
+  invocation_context.artifact_service = artifact_service
+  return ReadonlyContext(invocation_context)
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state():
+  instruction_template = "Hello {user_name}, you are in {app_state} state."
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo", "app_state": "active"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Hello Foo, you are in active state."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_without_placeholders_returns_template():
+  instruction_template = "A static instruction with no placeholders."
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == instruction_template
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_artifact():
+  instruction_template = "The artifact content is: {artifact.my_file}"
+  mock_artifact_service = MockArtifactService(
+      {"my_file": "This is my artifact content."}
+  )
+  invocation_context = await _create_test_readonly_context(
+      artifact_service=mock_artifact_service
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert (
+      populated_instruction
+      == "The artifact content is: This is my artifact content."
+  )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_optional_state():
+  instruction_template = "Optional value: {optional_value?}"
+  invocation_context = await _create_test_readonly_context()
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Optional value: "
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_missing_state_raises_key_error():
+  instruction_template = "Hello {missing_key}!"
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo"}
+  )
+
+  with pytest.raises(
+      KeyError, match="Context variable not found: `missing_key`."
+  ):
+    await instructions_utils.inject_session_state(
+        instruction_template, invocation_context
+    )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_missing_artifact_raises_key_error():
+  instruction_template = "The artifact content is: {artifact.missing_file}"
+  mock_artifact_service = MockArtifactService(
+      {"my_file": "This is my artifact content."}
+  )
+  invocation_context = await _create_test_readonly_context(
+      artifact_service=mock_artifact_service
+  )
+
+  with pytest.raises(KeyError, match="Artifact missing_file not found."):
+    await instructions_utils.inject_session_state(
+        instruction_template, invocation_context
+    )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_invalid_state_name_returns_original():
+  instruction_template = "Hello {invalid-key}!"
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Hello {invalid-key}!"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_invalid_prefix_state_name_returns_original():
+  instruction_template = "Hello {invalid:key}!"
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Hello {invalid:key}!"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_valid_prefix_state():
+  instruction_template = "Hello {app:user_name}!"
+  invocation_context = await _create_test_readonly_context(
+      state={"app:user_name": "Foo"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Hello Foo!"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_multiple_variables_and_artifacts():
+  instruction_template = """
+    Hello {user_name},
+    You are {user_age} years old.
+    Your favorite color is {favorite_color?}.
+    The artifact says: {artifact.my_file}
+    And another optional artifact: {artifact.other_file}
+    """
+  mock_artifact_service = MockArtifactService({
+      "my_file": "This is my artifact content.",
+      "other_file": "This is another artifact content.",
+  })
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo", "user_age": 30, "favorite_color": "blue"},
+      artifact_service=mock_artifact_service,
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  expected_instruction = """
+    Hello Foo,
+    You are 30 years old.
+    Your favorite color is blue.
+    The artifact says: This is my artifact content.
+    And another optional artifact: This is another artifact content.
+    """
+  assert populated_instruction == expected_instruction
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_empty_artifact_name_raises_key_error():
+  instruction_template = "The artifact content is: {artifact.}"
+  mock_artifact_service = MockArtifactService(
+      {"my_file": "This is my artifact content."}
+  )
+  invocation_context = await _create_test_readonly_context(
+      artifact_service=mock_artifact_service
+  )
+
+  with pytest.raises(KeyError, match="Artifact  not found."):
+    await instructions_utils.inject_session_state(
+        instruction_template, invocation_context
+    )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_artifact_service_not_initialized_raises_value_error():
+  instruction_template = "The artifact content is: {artifact.my_file}"
+  invocation_context = await _create_test_readonly_context()
+  with pytest.raises(ValueError, match="Artifact service is not initialized."):
+    await instructions_utils.inject_session_state(
+        instruction_template, invocation_context
+    )
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_optional_missing_artifact_returns_empty():
+  instruction_template = "Optional artifact: {artifact.missing_file?}"
+  mock_artifact_service = MockArtifactService(
+      {"my_file": "This is my artifact content."}
+  )
+  invocation_context = await _create_test_readonly_context(
+      artifact_service=mock_artifact_service
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Optional artifact: "
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_none_state_value_returns_empty():
+  instruction_template = "Value: {test_key}"
+  invocation_context = await _create_test_readonly_context(
+      state={"test_key": None}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Value: "
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_with_optional_missing_state_returns_empty():
+  instruction_template = "Optional value: {missing_key?}"
+  invocation_context = await _create_test_readonly_context()
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context
+  )
+  assert populated_instruction == "Optional value: "

tests/unittests/utils/test_model_name_utils.py

@@ -0,0 +1,447 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for model name utility functions."""
+
+from google.adk.models.llm_request import LlmRequest
+from google.adk.utils.model_name_utils import _is_gemini_3_x_live
+from google.adk.utils.model_name_utils import _is_managed_agent
+from google.adk.utils.model_name_utils import extract_model_name
+from google.adk.utils.model_name_utils import is_gemini_1_model
+from google.adk.utils.model_name_utils import is_gemini_3_5_live_translate
+from google.adk.utils.model_name_utils import is_gemini_eap_or_2_or_above
+from google.adk.utils.model_name_utils import is_gemini_model
+from google.adk.utils.model_name_utils import is_gemini_model_id_check_disabled
+
+
+class TestExtractModelName:
+  """Test the extract_model_name function."""
+
+  def test_extract_model_name_simple_model(self):
+    """Test extraction of simple model names."""
+    assert extract_model_name('gemini-2.5-pro') == 'gemini-2.5-pro'
+    assert extract_model_name('gemini-2.5-flash') == 'gemini-2.5-flash'
+    assert extract_model_name('gemini-1.0-pro') == 'gemini-1.0-pro'
+    assert extract_model_name('claude-3-sonnet') == 'claude-3-sonnet'
+    assert extract_model_name('gpt-4') == 'gpt-4'
+
+  def test_extract_model_name_path_based_model(self):
+    """Test extraction of path-based model names."""
+    path_model = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
+    assert extract_model_name(path_model) == 'gemini-2.5-flash'
+
+    path_model_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
+    assert extract_model_name(path_model_2) == 'gemini-1.5-pro-preview'
+
+    path_model_3 = 'projects/test-project/locations/europe-west1/publishers/google/models/claude-3-sonnet'
+    assert extract_model_name(path_model_3) == 'claude-3-sonnet'
+
+    path_model_4 = 'apigee/gemini-2.5-flash'
+    assert extract_model_name(path_model_4) == 'gemini-2.5-flash'
+
+    path_model_5 = 'apigee/v1/gemini-2.5-flash'
+    assert extract_model_name(path_model_5) == 'gemini-2.5-flash'
+
+    path_model_6 = 'apigee/gemini/gemini-2.5-flash'
+    assert extract_model_name(path_model_6) == 'gemini-2.5-flash'
+
+    path_model_7 = 'apigee/vertex_ai/gemini-2.5-flash'
+    assert extract_model_name(path_model_7) == 'gemini-2.5-flash'
+
+    path_model_8 = 'apigee/gemini/v1/gemini-2.5-flash'
+    assert extract_model_name(path_model_8) == 'gemini-2.5-flash'
+
+    path_model_9 = 'apigee/vertex_ai/v1beta/gemini-2.5-flash'
+    assert extract_model_name(path_model_9) == 'gemini-2.5-flash'
+
+  def test_extract_model_name_with_models_prefix(self):
+    """Test extraction of model names with 'models/' prefix."""
+    assert extract_model_name('models/gemini-2.5-pro') == 'gemini-2.5-pro'
+    assert extract_model_name('models/gemini-2.5-flash') == 'gemini-2.5-flash'
+
+  def test_extract_model_name_provider_prefixed_model(self):
+    """Test extraction of provider-prefixed Gemini model names."""
+    assert extract_model_name('gemini/gemini-2.5-flash') == 'gemini-2.5-flash'
+    assert extract_model_name('vertex_ai/gemini-2.5-flash') == (
+        'gemini-2.5-flash'
+    )
+    assert (
+        extract_model_name('openrouter/google/gemini-2.5-pro:online')
+        == 'gemini-2.5-pro:online'
+    )
+    assert extract_model_name('openrouter/anthropic/claude-sonnet-4') == (
+        'openrouter/anthropic/claude-sonnet-4'
+    )
+
+  def test_extract_model_name_invalid_path(self):
+    """Test that invalid path formats return the original string."""
+    invalid_paths = [
+        'projects/invalid/path/format',
+        'invalid/path/format',
+        'projects/123/locations/us-central1/models/gemini-2.5-flash',  # missing publishers
+        'projects/123/publishers/google/models/gemini-2.5-flash',  # missing locations
+        'projects/123/locations/us-central1/publishers/google/gemini-2.5-flash',  # missing models
+    ]
+
+    for invalid_path in invalid_paths:
+      assert extract_model_name(invalid_path) == invalid_path
+
+  def test_extract_model_name_empty_string(self):
+    """Test extraction from empty string."""
+    assert extract_model_name('') == ''
+
+  def test_extract_model_name_edge_cases(self):
+    """Test edge cases for model name extraction."""
+    # Test with unusual but valid path patterns
+    path_with_numbers = 'projects/123456789/locations/us-central1/publishers/google/models/gemini-2.5-flash'
+    assert extract_model_name(path_with_numbers) == 'gemini-2.5-flash'
+
+    # Test with hyphens in project/location names
+    path_with_hyphens = 'projects/my-test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro'
+    assert extract_model_name(path_with_hyphens) == 'gemini-1.5-pro'
+
+
+class TestIsGeminiModel:
+  """Test the is_gemini_model function."""
+
+  def test_is_gemini_model_simple_names(self):
+    """Test Gemini model detection with simple model names."""
+    assert is_gemini_model('gemini-2.5-pro') is True
+    assert is_gemini_model('gemini-1.5-flash') is True
+    assert is_gemini_model('gemini-1.0-pro') is True
+    assert is_gemini_model('gemini-2.5-flash') is True
+    assert is_gemini_model('claude-3-sonnet') is False
+    assert is_gemini_model('gpt-4') is False
+    assert is_gemini_model('llama-2') is False
+
+  def test_is_gemini_model_path_based_names(self):
+    """Test Gemini model detection with path-based model names."""
+    gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
+    assert is_gemini_model(gemini_path) is True
+
+    gemini_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.5-pro-preview'
+    assert is_gemini_model(gemini_path_2) is True
+
+    non_gemini_path = 'projects/265104255505/locations/us-central1/publishers/google/models/claude-3-sonnet'
+    assert is_gemini_model(non_gemini_path) is False
+
+  def test_is_gemini_model_provider_prefixed_names(self):
+    """Test Gemini model detection with provider-prefixed model names."""
+    assert is_gemini_model('gemini/gemini-2.5-flash') is True
+    assert is_gemini_model('vertex_ai/gemini-2.5-flash') is True
+    assert is_gemini_model('openrouter/google/gemini-2.5-pro:online') is True
+    assert is_gemini_model('openrouter/anthropic/claude-sonnet-4') is False
+
+  def test_is_gemini_model_edge_cases(self):
+    """Test edge cases for Gemini model detection."""
+    # Test with None
+    assert is_gemini_model(None) is False
+
+    # Test with empty string
+    assert is_gemini_model('') is False
+
+    # Test with model names containing gemini but not starting with it
+    assert is_gemini_model('my-gemini-model') is False
+    assert is_gemini_model('claude-gemini-hybrid') is False
+
+    # Test with model names that have gemini in the middle of the path
+    tricky_path = 'projects/265104255505/locations/us-central1/publishers/gemini/models/claude-3-sonnet'
+    assert is_gemini_model(tricky_path) is False
+
+    # Test with just "gemini" without dash
+    assert is_gemini_model('gemini') is False
+    assert is_gemini_model('gemini_1_5_flash') is False
+
+  def test_is_gemini_model_case_sensitivity(self):
+    """Test that model detection is case-sensitive."""
+    assert is_gemini_model('Gemini-2.5-pro') is False
+    assert is_gemini_model('GEMINI-2.5-pro') is False
+    assert is_gemini_model('gemini-2.5-PRO') is True  # Only the start matters
+
+
+class TestIsGemini1Model:
+  """Test the is_gemini_1_model function."""
+
+  def test_is_gemini_1_model_simple_names(self):
+    """Test Gemini 1.x model detection with simple model names."""
+    assert is_gemini_1_model('gemini-1.5-flash') is True
+    assert is_gemini_1_model('gemini-1.0-pro') is True
+    assert is_gemini_1_model('gemini-1.5-pro-preview') is True
+    assert is_gemini_1_model('gemini-1.9-experimental') is True
+    assert is_gemini_1_model('gemini-2.5-flash') is False
+    assert is_gemini_1_model('gemini-2.5-pro') is False
+    assert is_gemini_1_model('gemini-10.0-pro') is False  # Only 1.x versions
+    assert is_gemini_1_model('claude-3-sonnet') is False
+
+  def test_is_gemini_1_model_path_based_names(self):
+    """Test Gemini 1.x model detection with path-based model names."""
+    gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash'
+    assert is_gemini_1_model(gemini_1_path) is True
+
+    gemini_1_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-1.0-pro-preview'
+    assert is_gemini_1_model(gemini_1_path_2) is True
+
+    gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
+    assert is_gemini_1_model(gemini_2_path) is False
+
+  def test_is_gemini_1_model_provider_prefixed_names(self):
+    """Test Gemini 1.x detection with provider-prefixed model names."""
+    assert is_gemini_1_model('gemini/gemini-1.5-flash') is True
+    assert is_gemini_1_model('vertex_ai/gemini-1.5-flash') is True
+    assert is_gemini_1_model('openrouter/google/gemini-1.5-pro:online') is True
+    assert is_gemini_1_model('openrouter/google/gemini-2.5-pro') is False
+
+  def test_is_gemini_1_model_edge_cases(self):
+    """Test edge cases for Gemini 1.x model detection."""
+    # Test with None
+    assert is_gemini_1_model(None) is False
+
+    # Test with empty string
+    assert is_gemini_1_model('') is False
+
+    # Test with model names containing gemini-1 but not starting with it
+    assert is_gemini_1_model('my-gemini-1.5-model') is False
+    assert is_gemini_1_model('custom-gemini-1.5-flash') is False
+
+    # Test with invalid versions
+    assert is_gemini_1_model('gemini-1') is False  # Missing dot
+    assert is_gemini_1_model('gemini-1-pro') is False  # Missing dot
+    assert is_gemini_1_model('gemini-1.') is False  # Missing version number
+
+
+class TestIsGemini2Model:
+  """Test the is_gemini_eap_or_2_or_above function."""
+
+  def test_is_gemini_eap_or_2_or_above_simple_names(self):
+    """Test Gemini 2.0+ model detection with simple model names."""
+    assert is_gemini_eap_or_2_or_above('gemini-2.5-flash') is True
+    assert is_gemini_eap_or_2_or_above('gemini-2.5-pro') is True
+    assert is_gemini_eap_or_2_or_above('gemini-2.9-experimental') is True
+    assert is_gemini_eap_or_2_or_above('gemini-2-pro') is True
+    assert is_gemini_eap_or_2_or_above('gemini-2') is True
+    assert is_gemini_eap_or_2_or_above('gemini-3.0-pro') is True
+    assert is_gemini_eap_or_2_or_above('gemini-flash-early-exp') is True
+    assert is_gemini_eap_or_2_or_above('gemini-flash-early-exp3') is True
+    assert is_gemini_eap_or_2_or_above('gemini-flash-lite-early-exp') is True
+    assert is_gemini_eap_or_2_or_above('gemini-pro-early-exp') is True
+    assert is_gemini_eap_or_2_or_above('gemini-1.5-flash') is False
+    assert is_gemini_eap_or_2_or_above('gemini-1.0-pro') is False
+    assert is_gemini_eap_or_2_or_above('claude-3-sonnet') is False
+
+  def test_is_gemini_eap_or_2_or_above_path_based_names(self):
+    """Test Gemini 2.0+ model detection with path-based model names."""
+    gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.5-flash'
+    assert is_gemini_eap_or_2_or_above(gemini_2_path) is True
+
+    gemini_2_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-2.5-pro-preview'
+    assert is_gemini_eap_or_2_or_above(gemini_2_path_2) is True
+
+    gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash'
+    assert is_gemini_eap_or_2_or_above(gemini_1_path) is False
+
+    gemini_3_path = 'projects/12345/locations/us-east1/publishers/google/models/gemini-3.0-pro'
+    assert is_gemini_eap_or_2_or_above(gemini_3_path) is True
+
+  def test_is_gemini_eap_or_2_or_above_provider_prefixed_names(self):
+    """Test Gemini 2.0+ detection with provider-prefixed model names."""
+    assert is_gemini_eap_or_2_or_above('gemini/gemini-2.5-flash') is True
+    assert is_gemini_eap_or_2_or_above('vertex_ai/gemini-2.5-flash') is True
+    assert (
+        is_gemini_eap_or_2_or_above('openrouter/google/gemini-2.5-pro:online')
+        is True
+    )
+    assert (
+        is_gemini_eap_or_2_or_above('openrouter/google/gemini-1.5-pro:online')
+        is False
+    )
+
+  def test_is_gemini_eap_or_2_or_above_edge_cases(self):
+    """Test edge cases for Gemini 2.0+ model detection."""
+    # Test with None
+    assert is_gemini_eap_or_2_or_above(None) is False
+
+    # Test with empty string
+    assert is_gemini_eap_or_2_or_above('') is False
+
+    # Test with model names containing gemini-2 but not starting with it
+    assert is_gemini_eap_or_2_or_above('my-gemini-2.5-model') is False
+    assert is_gemini_eap_or_2_or_above('custom-gemini-2.5-flash') is False
+
+    # Test with invalid versions
+    assert (
+        is_gemini_eap_or_2_or_above('gemini-2.') is False
+    )  # Missing version number
+    assert is_gemini_eap_or_2_or_above('gemini-0.9-test') is False
+    assert is_gemini_eap_or_2_or_above('gemini-one') is False
+
+
+class TestModelNameUtilsIntegration:
+  """Integration tests for model name utilities."""
+
+  def test_model_classification_consistency(self):
+    """Test that model classification functions are consistent."""
+    test_models = [
+        'gemini-1.5-flash',
+        'gemini-2.5-flash',
+        'gemini-2.5-pro',
+        'gemini-3.0-pro',
+        'gemini/gemini-2.5-flash',
+        'openrouter/google/gemini-2.5-pro:online',
+        'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-pro',
+        'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-flash',
+        'projects/123/locations/us-central1/publishers/google/models/gemini-3.0-pro',
+        'claude-3-sonnet',
+        'gpt-4',
+    ]
+
+    for model in test_models:
+      # A model can only be either Gemini 1.x or Gemini 2.0+, not both
+      if is_gemini_1_model(model):
+        assert not is_gemini_eap_or_2_or_above(
+            model
+        ), f'Model {model} classified as both Gemini 1.x and 2.0+'
+        assert is_gemini_model(
+            model
+        ), f'Model {model} is Gemini 1.x but not classified as Gemini'
+
+      if is_gemini_eap_or_2_or_above(model):
+        assert not is_gemini_1_model(
+            model
+        ), f'Model {model} classified as both Gemini 1.x and 2.0+'
+        assert is_gemini_model(
+            model
+        ), f'Model {model} is Gemini 2.0+ but not classified as Gemini'
+
+      # If it's neither Gemini 1.x nor 2.0+, it should not be classified as Gemini
+      if not is_gemini_1_model(model) and not is_gemini_eap_or_2_or_above(
+          model
+      ):
+        if model and 'gemini-' not in extract_model_name(model):
+          assert not is_gemini_model(
+              model
+          ), f'Non-Gemini model {model} classified as Gemini'
+
+  def test_path_vs_simple_model_consistency(self):
+    """Test that path-based and simple model names are classified consistently."""
+    model_pairs = [
+        (
+            'gemini-1.5-flash',
+            'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-flash',
+        ),
+        (
+            'gemini-2.5-flash',
+            'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-flash',
+        ),
+        (
+            'gemini-2.5-pro',
+            'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-pro',
+        ),
+        (
+            'gemini-3.0-pro',
+            'projects/123/locations/us-central1/publishers/google/models/gemini-3.0-pro',
+        ),
+        (
+            'claude-3-sonnet',
+            'projects/123/locations/us-central1/publishers/google/models/claude-3-sonnet',
+        ),
+    ]
+
+    for simple_model, path_model in model_pairs:
+      # Both forms should be classified identically
+      assert is_gemini_model(simple_model) == is_gemini_model(path_model), (
+          f'Inconsistent Gemini classification for {simple_model} vs'
+          f' {path_model}'
+      )
+      assert is_gemini_1_model(simple_model) == is_gemini_1_model(path_model), (
+          f'Inconsistent Gemini 1.x classification for {simple_model} vs'
+          f' {path_model}'
+      )
+      assert is_gemini_eap_or_2_or_above(
+          simple_model
+      ) == is_gemini_eap_or_2_or_above(path_model), (
+          f'Inconsistent Gemini 2.0+ classification for {simple_model} vs'
+          f' {path_model}'
+      )
+
+
+class TestGeminiModelIdCheckFlag:
+  """Tests for Gemini model-id check override flag."""
+
+  def test_default_is_disabled(self, monkeypatch):
+    monkeypatch.delenv('ADK_DISABLE_GEMINI_MODEL_ID_CHECK', raising=False)
+    assert is_gemini_model_id_check_disabled() is False
+
+  def test_true_enables_check_bypass(self, monkeypatch):
+    monkeypatch.setenv('ADK_DISABLE_GEMINI_MODEL_ID_CHECK', 'true')
+    assert is_gemini_model_id_check_disabled() is True
+
+
+class TestIsGemini3XLive:
+  """Test the _is_gemini_3_x_live function."""
+
+  def test_is_gemini_3_x_live_simple_name(self):
+    """Test with simple model name format."""
+    assert _is_gemini_3_x_live('gemini-3.1-flash-live') is True
+    assert _is_gemini_3_x_live('gemini-3.1-flash-live-preview') is True
+    assert _is_gemini_3_x_live('gemini-3.5-flash-lite-live-preview') is True
+    assert _is_gemini_3_x_live('gemini-3.5-live-translate') is False
+    assert _is_gemini_3_x_live('gemini-3.1-pro') is False
+    assert _is_gemini_3_x_live('gemini-2.5-flash-live') is False
+
+  def test_is_gemini_3_x_live_path_based_name(self):
+    """Test with path-based format (Vertex AI etc.)."""
+    vertex_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash-live'
+    assert _is_gemini_3_x_live(vertex_path) is True
+
+    vertex_path_preview = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite-live-preview'
+    assert _is_gemini_3_x_live(vertex_path_preview) is True
+
+    non_live_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash'
+    assert _is_gemini_3_x_live(non_live_path) is False
+
+  def test_is_gemini_3_x_live_edge_cases(self):
+    """Test edge cases."""
+    assert _is_gemini_3_x_live(None) is False
+    assert _is_gemini_3_x_live('') is False
+
+
+class TestIsGemini35LiveTranslate:
+  """Test the is_gemini_3_5_live_translate function."""
+
+  def test_is_gemini_3_5_live_translate_simple_name(self):
+    """Test with simple model name format."""
+    assert is_gemini_3_5_live_translate('gemini-3.5-live-translate') is True
+    assert is_gemini_3_5_live_translate('gemini-3.5-flash-live') is False
+
+  def test_is_gemini_3_5_live_translate_path_based_name(self):
+    """Test with path-based format (Vertex AI etc.)."""
+    vertex_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.5-live-translate-preview'
+    assert is_gemini_3_5_live_translate(vertex_path) is True
+
+  def test_is_gemini_3_5_live_translate_edge_cases(self):
+    """Test edge cases."""
+    assert is_gemini_3_5_live_translate(None) is False
+    assert is_gemini_3_5_live_translate('') is False
+
+
+class TestIsManagedAgent:
+  """Tests for the _is_managed_agent predicate."""
+
+  def test_true_when_flag_set(self):
+    request = LlmRequest()
+    request._is_managed_agent = True
+    assert _is_managed_agent(request) is True
+
+  def test_false_by_default(self):
+    assert _is_managed_agent(LlmRequest()) is False

tests/unittests/utils/test_mtls_utils.py

@@ -0,0 +1,303 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for _mtls_utils."""
+
+import os
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.utils import _mtls_utils
+from google.auth import exceptions as ga_exceptions
+from google.auth.transport import mtls
+import pytest
+
+_DEFAULT_TEMPLATE = "service.{location}.rep.googleapis.com"
+_MTLS_TEMPLATE = "service.{location}.rep.mtls.googleapis.com"
+_LOCATION = "us-central1"
+
+
+class TestMtlsUtils:
+  """Tests for _mtls_utils functions."""
+
+  @patch.object(mtls, "should_use_client_cert", autospec=True)
+  def test_use_client_cert_effective_with_mtls_cert_true(
+      self, mock_should_use_client_cert
+  ):
+    mock_should_use_client_cert.return_value = True
+    assert _mtls_utils.use_client_cert_effective() is True
+    mock_should_use_client_cert.assert_called_once()
+
+  @patch.object(mtls, "should_use_client_cert", autospec=True)
+  def test_use_client_cert_effective_with_mtls_cert_false(
+      self, mock_should_use_client_cert
+  ):
+    mock_should_use_client_cert.return_value = False
+    assert _mtls_utils.use_client_cert_effective() is False
+    mock_should_use_client_cert.assert_called_once()
+
+  @patch.object(mtls, "should_use_client_cert", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"})
+  def test_use_client_cert_effective_fallback_true(
+      self, mock_should_use_client_cert
+  ):
+    mock_should_use_client_cert.side_effect = AttributeError
+    assert _mtls_utils.use_client_cert_effective() is True
+
+  @patch.object(mtls, "should_use_client_cert", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"})
+  def test_use_client_cert_effective_fallback_false(
+      self, mock_should_use_client_cert
+  ):
+    mock_should_use_client_cert.side_effect = AttributeError
+    assert _mtls_utils.use_client_cert_effective() is False
+
+  @patch.object(mtls, "should_use_client_cert", autospec=True)
+  @patch.dict("os.environ", {}, clear=True)
+  def test_use_client_cert_effective_fallback_default_false(
+      self, mock_should_use_client_cert
+  ):
+    mock_should_use_client_cert.side_effect = AttributeError
+    assert _mtls_utils.use_client_cert_effective() is False
+
+  @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"})
+  def test_get_api_endpoint_always(self, mock_use_client_cert):
+    endpoint = _mtls_utils.get_api_endpoint(
+        _LOCATION, _DEFAULT_TEMPLATE, _MTLS_TEMPLATE
+    )
+    assert endpoint == _MTLS_TEMPLATE.format(location=_LOCATION)
+    mock_use_client_cert.assert_not_called()
+
+  @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"})
+  def test_get_api_endpoint_never(self, mock_use_client_cert):
+    endpoint = _mtls_utils.get_api_endpoint(
+        _LOCATION, _DEFAULT_TEMPLATE, _MTLS_TEMPLATE
+    )
+    assert endpoint == _DEFAULT_TEMPLATE.format(location=_LOCATION)
+    mock_use_client_cert.assert_not_called()
+
+  @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"})
+  def test_get_api_endpoint_auto_with_cert(self, mock_use_client_cert):
+    mock_use_client_cert.return_value = True
+    endpoint = _mtls_utils.get_api_endpoint(
+        _LOCATION, _DEFAULT_TEMPLATE, _MTLS_TEMPLATE
+    )
+    assert endpoint == _MTLS_TEMPLATE.format(location=_LOCATION)
+    mock_use_client_cert.assert_called_once()
+
+  @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"})
+  def test_get_api_endpoint_auto_without_cert(self, mock_use_client_cert):
+    mock_use_client_cert.return_value = False
+    endpoint = _mtls_utils.get_api_endpoint(
+        _LOCATION, _DEFAULT_TEMPLATE, _MTLS_TEMPLATE
+    )
+    assert endpoint == _DEFAULT_TEMPLATE.format(location=_LOCATION)
+    mock_use_client_cert.assert_called_once()
+
+  @patch.object(_mtls_utils, "use_client_cert_effective", autospec=True)
+  @patch.dict("os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "invalid_value"})
+  def test_get_api_endpoint_invalid_fallback_to_auto(
+      self, mock_use_client_cert
+  ):
+    mock_use_client_cert.return_value = True
+    endpoint = _mtls_utils.get_api_endpoint(
+        _LOCATION, _DEFAULT_TEMPLATE, _MTLS_TEMPLATE
+    )
+    assert endpoint == _MTLS_TEMPLATE.format(location=_LOCATION)
+    mock_use_client_cert.assert_called_once()
+
+  @pytest.mark.parametrize(
+      "url, expected",
+      [
+          ("https://oauth2.googleapis.com/token", True),
+          ("https://openidconnect.googleapis.com/v1/userinfo", True),
+          ("https://oauth2.mtls.googleapis.com/token", False),
+          ("https://example.com/token", False),
+          ("https://accounts.google.com/o/oauth2/v2/auth", False),
+          (None, False),
+          ("", False),
+      ],
+  )
+  def test_is_non_mtls_googleapis_endpoint(self, url, expected):
+    assert _mtls_utils.is_non_mtls_googleapis_endpoint(url) is expected
+
+  @patch.dict("os.environ", {}, clear=True)
+  @pytest.mark.parametrize(
+      "url, expected",
+      [
+          (
+              "https://oauth2.googleapis.com/token",
+              "https://oauth2.mtls.googleapis.com/token",
+          ),
+          (
+              "https://openidconnect.googleapis.com/v1/userinfo",
+              "https://openidconnect.mtls.googleapis.com/v1/userinfo",
+          ),
+          # Non-Google providers are never rewritten.
+          ("https://example.com/token", "https://example.com/token"),
+          # Already-mTLS hosts are left alone.
+          (
+              "https://oauth2.mtls.googleapis.com/token",
+              "https://oauth2.mtls.googleapis.com/token",
+          ),
+      ],
+  )
+  def test_effective_googleapis_endpoint_rewrites(self, url, expected):
+    assert _mtls_utils.effective_googleapis_endpoint(url) == expected
+
+  @patch.dict(
+      "os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}, clear=True
+  )
+  def test_effective_googleapis_endpoint_never_opts_out(self):
+    url = "https://oauth2.googleapis.com/token"
+    assert _mtls_utils.effective_googleapis_endpoint(url) == url
+
+  @patch.dict("os.environ", {}, clear=True)
+  def test_effective_googleapis_endpoint_preserves_query(self):
+    url = "https://iam.googleapis.com/v1/token?foo=bar"
+    assert (
+        _mtls_utils.effective_googleapis_endpoint(url)
+        == "https://iam.mtls.googleapis.com/v1/token?foo=bar"
+    )
+
+  @patch("google.auth.transport.mtls.has_default_client_cert_source")
+  @patch("google.auth.transport._mtls_helper.get_client_cert_and_key")
+  @patch("google.auth.transport.requests._MutualTlsAdapter")
+  def test_configure_session_for_mtls_mounts_adapter(
+      self, mock_adapter, mock_get_cert, mock_has_cert_source
+  ):
+    mock_has_cert_source.return_value = False
+    mock_get_cert.return_value = (True, b"cert", b"key")
+    session = MagicMock()
+
+    result = _mtls_utils.configure_session_for_mtls(session)
+
+    assert result is True
+    mock_adapter.assert_called_once_with(b"cert", b"key")
+    session.mount.assert_called_once_with("https://", mock_adapter.return_value)
+
+  @patch("google.auth.transport.mtls.has_default_client_cert_source")
+  @patch("google.auth.transport._mtls_helper.get_client_cert_and_key")
+  def test_configure_session_for_mtls_no_cert(
+      self, mock_get_cert, mock_has_cert_source
+  ):
+    mock_has_cert_source.return_value = False
+    mock_get_cert.return_value = (False, None, None)
+    session = MagicMock()
+
+    result = _mtls_utils.configure_session_for_mtls(session)
+
+    assert result is False
+    session.mount.assert_not_called()
+
+  @patch("google.auth.transport.mtls.has_default_client_cert_source")
+  @patch("google.auth.transport._mtls_helper.get_client_cert_and_key")
+  def test_configure_session_for_mtls_cert_error_falls_back(
+      self, mock_get_cert, mock_has_cert_source
+  ):
+    mock_has_cert_source.return_value = False
+    mock_get_cert.side_effect = ga_exceptions.ClientCertError("boom")
+    session = MagicMock()
+
+    result = _mtls_utils.configure_session_for_mtls(session)
+
+    assert result is False
+    session.mount.assert_not_called()
+
+
+class TestMtlsClientCerts:
+  """Tests for MtlsClientCerts."""
+
+  @patch.object(mtls, "has_default_client_cert_source", autospec=True)
+  def test_get_certs_no_default_source(self, mock_has_cert):
+    mock_has_cert.return_value = False
+    certs = _mtls_utils.MtlsClientCerts()
+    cert_path, key_path, passphrase = certs.get_certs()
+    assert cert_path is None
+    assert key_path is None
+    assert passphrase is None
+    mock_has_cert.assert_called_once()
+
+  @patch.object(mtls, "has_default_client_cert_source", autospec=True)
+  @patch.object(mtls, "default_client_encrypted_cert_source", autospec=True)
+  def test_get_certs_with_default_source(
+      self, mock_encrypted_source, mock_has_cert
+  ):
+    mock_has_cert.return_value = True
+
+    mock_cert_source = MagicMock()
+    mock_cert_source.return_value = (None, None, b"test_passphrase")
+    mock_encrypted_source.return_value = mock_cert_source
+
+    certs = _mtls_utils.MtlsClientCerts()
+    cert_path, key_path, passphrase = certs.get_certs()
+
+    assert cert_path is not None
+    assert key_path is not None
+    assert passphrase == b"test_passphrase"
+
+    assert os.path.exists(certs._tempdir.name)
+    assert cert_path.startswith(certs._tempdir.name)
+    assert key_path.startswith(certs._tempdir.name)
+
+    # Getting certs again should return cached values without calling mtls again
+    cert_path2, key_path2, passphrase2 = certs.get_certs()
+    assert cert_path2 == cert_path
+    assert key_path2 == key_path
+    assert passphrase2 == passphrase
+    mock_has_cert.assert_called_once()
+    mock_encrypted_source.assert_called_once()
+
+  @patch.object(mtls, "has_default_client_cert_source", autospec=True)
+  @patch.object(mtls, "default_client_encrypted_cert_source", autospec=True)
+  def test_get_certs_extraction_failure(
+      self, mock_encrypted_source, mock_has_cert
+  ):
+    mock_has_cert.return_value = True
+    mock_encrypted_source.side_effect = Exception("extraction failed")
+
+    certs = _mtls_utils.MtlsClientCerts()
+    with pytest.raises(
+        RuntimeError, match="Failed to extract default client certificates"
+    ):
+      certs.get_certs()
+
+    assert certs._tempdir is None
+
+  @patch.object(mtls, "has_default_client_cert_source", autospec=True)
+  @patch.object(mtls, "default_client_encrypted_cert_source", autospec=True)
+  def test_close_cleans_up_tempdir(self, mock_encrypted_source, mock_has_cert):
+    mock_has_cert.return_value = True
+    mock_cert_source = MagicMock()
+    mock_cert_source.return_value = (None, None, b"test_passphrase")
+    mock_encrypted_source.return_value = mock_cert_source
+
+    certs = _mtls_utils.MtlsClientCerts()
+    certs.get_certs()
+
+    tempdir_name = certs._tempdir.name
+    assert os.path.exists(tempdir_name)
+
+    certs.close()
+
+    assert not os.path.exists(tempdir_name)
+    assert certs._tempdir is None
+    assert certs.cert_path is None
+    assert certs.key_path is None
+    assert certs.passphrase is None
+    assert not certs._initialized

tests/unittests/utils/test_output_schema_utils.py

@@ -0,0 +1,119 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import importlib.util
+
+from google.adk.models.base_llm import BaseLlm
+from google.adk.models.google_llm import Gemini
+from google.adk.utils.output_schema_utils import can_use_output_schema_with_tools
+import pytest
+
+_has_anthropic = importlib.util.find_spec("anthropic") is not None
+_has_litellm = importlib.util.find_spec("litellm") is not None
+
+_skip_anthropic = pytest.mark.skipif(
+    not _has_anthropic, reason="anthropic not installed"
+)
+_skip_litellm = pytest.mark.skipif(
+    not _has_litellm, reason="litellm not installed"
+)
+
+
+def _make_claude(model: str):
+  from google.adk.models.anthropic_llm import Claude
+
+  return Claude(model=model)
+
+
+def _make_litellm(model: str):
+  from google.adk.models.lite_llm import LiteLlm
+
+  return LiteLlm(model=model)
+
+
+@pytest.mark.parametrize(
+    "model, env_value, expected",
+    [
+        ("gemini-2.5-pro", "1", True),
+        ("gemini-2.5-pro", "0", False),
+        ("gemini-2.5-pro", None, False),
+        (Gemini(model="gemini-2.5-pro"), "1", True),
+        (Gemini(model="gemini-2.5-pro"), "0", False),
+        (Gemini(model="gemini-2.5-pro"), None, False),
+        ("gemini-2.5-flash", "1", True),
+        ("gemini-2.5-flash", "0", False),
+        ("gemini-2.5-flash", None, False),
+        ("gemini-1.5-pro", "1", False),
+        ("gemini-1.5-pro", "0", False),
+        ("gemini-1.5-pro", None, False),
+    ],
+)
+def test_can_use_output_schema_with_tools(
+    monkeypatch: pytest.MonkeyPatch,
+    model: str | BaseLlm,
+    env_value: str | None,
+    expected: bool,
+) -> None:
+  """Test can_use_output_schema_with_tools."""
+  if env_value is not None:
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", env_value)
+  else:
+    monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False)
+  assert can_use_output_schema_with_tools(model) == expected
+
+
+@_skip_anthropic
+@pytest.mark.parametrize(
+    "model, env_value, expected",
+    [
+        ("claude-3.7-sonnet", "1", False),
+        ("claude-3.7-sonnet", "0", False),
+        ("claude-3.7-sonnet", None, False),
+    ],
+)
+def test_can_use_output_schema_with_tools_claude(
+    monkeypatch, model, env_value, expected
+):
+  """Test can_use_output_schema_with_tools with Claude models."""
+  claude_model = _make_claude(model)
+  if env_value is not None:
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", env_value)
+  else:
+    monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False)
+  assert can_use_output_schema_with_tools(claude_model) == expected
+
+
+@_skip_litellm
+@pytest.mark.parametrize(
+    "model, env_value, expected",
+    [
+        ("openai/gpt-4o", "1", True),
+        ("openai/gpt-4o", "0", True),
+        ("openai/gpt-4o", None, True),
+        ("anthropic/claude-3.7-sonnet", None, True),
+        ("fireworks_ai/llama-v3p1-70b", None, True),
+    ],
+)
+def test_can_use_output_schema_with_tools_litellm(
+    monkeypatch, model, env_value, expected
+):
+  """Test can_use_output_schema_with_tools with LiteLLM models."""
+  litellm_model = _make_litellm(model)
+  if env_value is not None:
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", env_value)
+  else:
+    monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False)
+  assert can_use_output_schema_with_tools(litellm_model) == expected

tests/unittests/utils/test_schema_utils.py

@@ -0,0 +1,210 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for _schema_utils module."""
+
+from google.adk.utils._schema_utils import get_list_inner_type
+from google.adk.utils._schema_utils import is_basemodel_schema
+from google.adk.utils._schema_utils import is_list_of_basemodel
+from google.adk.utils._schema_utils import validate_node_data
+from google.adk.utils._schema_utils import validate_schema
+from google.genai import types
+from pydantic import BaseModel
+from pydantic import ValidationError
+import pytest
+
+
+class SampleModel(BaseModel):
+  """Sample model for testing."""
+
+  name: str
+  value: int
+
+
+class TestIsBasemodelSchema:
+  """Tests for is_basemodel_schema function."""
+
+  def test_basemodel_class_returns_true(self):
+    """Test that a BaseModel class returns True."""
+    assert is_basemodel_schema(SampleModel)
+
+  def test_list_of_basemodel_returns_false(self):
+    """Test that list[BaseModel] returns False."""
+    assert not is_basemodel_schema(list[SampleModel])
+
+  def test_list_of_str_returns_false(self):
+    """Test that list[str] returns False."""
+    assert not is_basemodel_schema(list[str])
+
+  def test_dict_returns_false(self):
+    """Test that dict types return False."""
+    assert not is_basemodel_schema(dict[str, int])
+
+  def test_plain_str_returns_false(self):
+    """Test that plain str returns False."""
+    assert not is_basemodel_schema(str)
+
+  def test_plain_int_returns_false(self):
+    """Test that plain int returns False."""
+    assert not is_basemodel_schema(int)
+
+
+class TestIsListOfBasemodel:
+  """Tests for is_list_of_basemodel function."""
+
+  def test_list_of_basemodel_returns_true(self):
+    """Test that list[BaseModel] returns True."""
+    assert is_list_of_basemodel(list[SampleModel])
+
+  def test_basemodel_class_returns_false(self):
+    """Test that a plain BaseModel class returns False."""
+    assert not is_list_of_basemodel(SampleModel)
+
+  def test_list_of_str_returns_false(self):
+    """Test that list[str] returns False."""
+    assert not is_list_of_basemodel(list[str])
+
+  def test_list_of_int_returns_false(self):
+    """Test that list[int] returns False."""
+    assert not is_list_of_basemodel(list[int])
+
+  def test_dict_returns_false(self):
+    """Test that dict types return False."""
+    assert not is_list_of_basemodel(dict[str, int])
+
+  def test_plain_list_returns_false(self):
+    """Test that plain list (no type arg) returns False."""
+    assert not is_list_of_basemodel(list)
+
+
+class TestGetListInnerType:
+  """Tests for get_list_inner_type function."""
+
+  def test_list_of_basemodel_returns_inner_type(self):
+    """Test that list[BaseModel] returns the inner type."""
+    assert get_list_inner_type(list[SampleModel]) is SampleModel
+
+  def test_basemodel_class_returns_none(self):
+    """Test that a plain BaseModel class returns None."""
+    assert get_list_inner_type(SampleModel) is None
+
+  def test_list_of_str_returns_none(self):
+    """Test that list[str] returns None."""
+    assert get_list_inner_type(list[str]) is None
+
+  def test_dict_returns_none(self):
+    """Test that dict types return None."""
+    assert get_list_inner_type(dict[str, int]) is None
+
+
+class TestValidateSchema:
+  """Tests for validate_schema function."""
+
+  def test_basemodel_schema(self):
+    """Test validation with a BaseModel schema."""
+    json_text = '{"name": "test", "value": 42}'
+    result = validate_schema(SampleModel, json_text)
+    assert result == {'name': 'test', 'value': 42}
+
+  def test_basemodel_schema_excludes_none(self):
+    """Test that None values are excluded from the result."""
+
+    class ModelWithOptional(BaseModel):
+      name: str
+      optional_field: str | None = None
+
+    json_text = '{"name": "test", "optional_field": null}'
+    result = validate_schema(ModelWithOptional, json_text)
+    assert result == {'name': 'test'}
+
+  def test_list_of_basemodel_schema(self):
+    """Test validation with a list[BaseModel] schema."""
+    json_text = '[{"name": "item1", "value": 1}, {"name": "item2", "value": 2}]'
+    result = validate_schema(list[SampleModel], json_text)
+    assert result == [
+        {'name': 'item1', 'value': 1},
+        {'name': 'item2', 'value': 2},
+    ]
+
+  def test_list_of_str_schema(self):
+    """Test validation with a list[str] schema."""
+    json_text = '["a", "b", "c"]'
+    result = validate_schema(list[str], json_text)
+    assert result == ['a', 'b', 'c']
+
+  def test_dict_schema(self):
+    """Test validation with a dict schema."""
+    json_text = '{"key1": 1, "key2": 2}'
+    result = validate_schema(dict[str, int], json_text)
+    assert result == {'key1': 1, 'key2': 2}
+
+
+class TestValidateNodeData:
+  """Tests for validate_node_data function."""
+
+  def test_none_schema_or_data_returns_data(self):
+    """Bypasses validation if schema or data is None."""
+    assert validate_node_data(None, 'some_data') == 'some_data'
+    assert validate_node_data(SampleModel, None) is None
+
+  def test_dict_or_types_schema_returns_data(self):
+    """Bypasses validation if schema is dict or types.Schema."""
+    assert validate_node_data({'key': int}, 'some_data') == 'some_data'
+    # Mock types.Schema
+    schema = types.Schema(type=types.Type.STRING)
+    assert validate_node_data(schema, 'some_data') == 'some_data'
+
+  def test_content_schema_returns_data(self):
+    """Bypasses validation if target schema is types.Content or subclass."""
+    result = validate_node_data(
+        types.Content, types.Content(role='user', parts=[])
+    )
+    assert result == {'role': 'user', 'parts': []}
+
+  def test_plain_basemodel_schema_validates_raw_dict(self):
+    """Validates raw dict data against BaseModel schema."""
+    result = validate_node_data(SampleModel, {'name': 'test', 'value': 42})
+    assert result == {'name': 'test', 'value': 42}
+
+  def test_content_data_and_preserve_content(self):
+    """Validates wrapped content and wraps result back into Content."""
+    data = types.Content(
+        role='user',
+        parts=[types.Part(text='{"name": "test", "value": 42}')],
+    )
+    result = validate_node_data(SampleModel, data, preserve_content=True)
+    assert isinstance(result, types.Content)
+    assert result.role == 'user'
+    assert len(result.parts) == 1
+    assert result.parts[0].text == '{"name": "test", "value": 42}'
+
+  def test_content_data_no_preserve_content(self):
+    """Validates wrapped content and returns unwrapped dictionary."""
+    data = types.Content(
+        role='user',
+        parts=[types.Part(text='{"name": "test", "value": 42}')],
+    )
+    result = validate_node_data(SampleModel, data, preserve_content=False)
+    assert isinstance(result, dict)
+    assert result == {'name': 'test', 'value': 42}
+
+  def test_raw_json_string_validated_against_basemodel_schema(self):
+    """Raw JSON string fails validation against BaseModel schema (not auto-parsed)."""
+    with pytest.raises(ValidationError):
+      validate_node_data(SampleModel, '{"name": "test", "value": 42}')
+
+  def test_raw_string_not_parsed_with_str_schema(self):
+    """Bypasses JSON parsing if schema is str."""
+    result = validate_node_data(str, 'hello')
+    assert result == 'hello'

tests/unittests/utils/test_serialized_base_model.py

@@ -0,0 +1,35 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for SerializedBaseModel."""
+
+from google.adk.utils._serialized_base_model import SerializedBaseModel
+
+
+class MyModel(SerializedBaseModel):
+  test_field: str
+
+
+def test_model_dump_json_by_alias_default():
+  model = MyModel(test_field="value")
+  json_str = model.model_dump_json()
+  assert "testField" in json_str
+  assert "test_field" not in json_str
+
+
+def test_model_dump_json_by_alias_false():
+  model = MyModel(test_field="value")
+  json_str = model.model_dump_json(by_alias=False)
+  assert "test_field" in json_str
+  assert "testField" not in json_str

tests/unittests/utils/test_streaming_utils.py

@@ -0,0 +1,718 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.features._feature_registry import FeatureName
+from google.adk.features._feature_registry import temporary_feature_override
+from google.adk.flows.llm_flows.functions import AF_FUNCTION_CALL_ID_PREFIX
+from google.adk.utils import streaming_utils
+from google.genai import types
+import pytest
+
+
+class TestStreamingResponseAggregator:
+
+  @pytest.mark.asyncio
+  async def test_process_response_with_text(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="Hello")])
+            )
+        ]
+    )
+    results = []
+    async for r in aggregator.process_response(response):
+      results.append(r)
+    assert len(results) == 1
+    assert results[0].content.parts[0].text == "Hello"
+    assert results[0].partial
+
+  @pytest.mark.asyncio
+  async def test_process_response_with_thought(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(
+                    parts=[types.Part(text="Thinking...", thought=True)]
+                )
+            )
+        ]
+    )
+    results = []
+    async for r in aggregator.process_response(response):
+      results.append(r)
+    assert len(results) == 1
+    assert results[0].content.parts[0].text == "Thinking..."
+    assert results[0].content.parts[0].thought
+    assert results[0].partial
+
+  @pytest.mark.asyncio
+  async def test_process_response_multiple(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response1 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="Hello ")])
+            )
+        ]
+    )
+    response2 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="World!")])
+            )
+        ]
+    )
+    async for _ in aggregator.process_response(response1):
+      pass
+    results = []
+    async for r in aggregator.process_response(response2):
+      results.append(r)
+    assert len(results) == 1
+    assert results[0].content.parts[0].text == "World!"
+
+    closed_response = aggregator.close()
+    assert closed_response is not None
+    assert closed_response.content.parts[0].text == "Hello World!"
+
+  @pytest.mark.asyncio
+  async def test_process_response_interleaved_thought_and_text(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response1 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(
+                    parts=[types.Part(text="I am thinking...", thought=True)]
+                )
+            )
+        ]
+    )
+    response2 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(
+                    parts=[types.Part(text="Okay, I have a result.")]
+                )
+            )
+        ]
+    )
+    response3 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(
+                    parts=[types.Part(text=" The result is 42.")]
+                )
+            )
+        ]
+    )
+
+    async for _ in aggregator.process_response(response1):
+      pass
+    async for _ in aggregator.process_response(response2):
+      pass
+    async for _ in aggregator.process_response(response3):
+      pass
+
+    closed_response = aggregator.close()
+    assert closed_response is not None
+    assert len(closed_response.content.parts) == 2
+    assert closed_response.content.parts[0].text == "I am thinking..."
+    assert closed_response.content.parts[0].thought
+    assert (
+        closed_response.content.parts[1].text
+        == "Okay, I have a result. The result is 42."
+    )
+    assert not closed_response.content.parts[1].thought
+
+  def test_close_with_no_responses(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    closed_response = aggregator.close()
+    assert closed_response is None
+
+  @pytest.mark.asyncio
+  async def test_close_with_finish_reason(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="Hello")]),
+                finish_reason=types.FinishReason.STOP,
+            )
+        ]
+    )
+    async for _ in aggregator.process_response(response):
+      pass
+    closed_response = aggregator.close()
+    assert closed_response is not None
+    assert closed_response.content.parts[0].text == "Hello"
+    assert closed_response.error_code is None
+    assert closed_response.error_message is None
+
+  @pytest.mark.asyncio
+  async def test_close_with_error(self):
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="Error")]),
+                finish_reason=types.FinishReason.RECITATION,
+                finish_message="Recitation error",
+            )
+        ]
+    )
+    async for _ in aggregator.process_response(response):
+      pass
+    closed_response = aggregator.close()
+    assert closed_response is not None
+    assert closed_response.content.parts[0].text == "Error"
+    assert closed_response.error_code == types.FinishReason.RECITATION
+    assert closed_response.error_message == "Recitation error"
+
+  @pytest.mark.asyncio
+  @pytest.mark.parametrize("use_progressive_sse", [True, False])
+  async def test_empty_content_produces_empty_final_frame(
+      self, use_progressive_sse
+  ):
+    """A candidate with empty parts + STOP passes through without an error.
+
+    A terminal empty STOP chunk must not be classified as an error at the
+    streaming layer; consumers that batch parts across chunks rely on it
+    passing through cleanly.
+    """
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+      response = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[]),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ]
+      )
+      results = []
+      async for r in aggregator.process_response(response):
+        results.append(r)
+      closed_response = aggregator.close()
+
+      assert len(results) == 1
+      assert results[0].content is not None
+      assert results[0].error_code is None
+      assert closed_response is not None
+      assert closed_response.partial is False
+      assert closed_response.content is None
+      assert closed_response.finish_reason == types.FinishReason.STOP
+
+  @pytest.mark.asyncio
+  @pytest.mark.parametrize("use_progressive_sse", [True, False])
+  async def test_prompt_feedback_block_returns_error_frame(
+      self, use_progressive_sse
+  ):
+    """A prompt-level safety block produces a final frame with the error code."""
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+      response = types.GenerateContentResponse(
+          prompt_feedback=types.GenerateContentResponsePromptFeedback(
+              block_reason=types.BlockedReason.SAFETY,
+              block_reason_message="Blocked by safety",
+          )
+      )
+      results = []
+      async for r in aggregator.process_response(response):
+        results.append(r)
+      closed_response = aggregator.close()
+
+      assert len(results) == 1
+      assert closed_response is not None
+      assert closed_response.partial is False
+      assert closed_response.error_code == types.BlockedReason.SAFETY
+      assert closed_response.error_message == "Blocked by safety"
+      assert closed_response.content is None
+
+  @pytest.mark.asyncio
+  @pytest.mark.parametrize("use_progressive_sse", [True, False])
+  async def test_pure_function_call_behavior_differs_by_mode(
+      self, use_progressive_sse
+  ):
+    """A pure function call yields the part in progressive mode and an empty frame otherwise."""
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, use_progressive_sse
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+      response = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name="my_tool",
+                                  args={"x": 1},
+                              )
+                          )
+                      ]
+                  ),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ]
+      )
+
+      results = []
+      async for r in aggregator.process_response(response):
+        results.append(r)
+      closed_response = aggregator.close()
+
+      assert closed_response is not None
+      assert closed_response.partial is False
+
+      if use_progressive_sse:
+        assert closed_response.content is not None
+        assert len(closed_response.content.parts) == 1
+        assert closed_response.content.parts[0].function_call.name == "my_tool"
+      else:
+        assert closed_response.content is None
+
+  @pytest.mark.asyncio
+  @pytest.mark.parametrize(
+      "test_id, use_progressive_sse, metadata_type",
+      [
+          ("grounding_default", False, "grounding"),
+          ("grounding_progressive", True, "grounding"),
+          ("citation_default", False, "citation"),
+          ("citation_progressive", True, "citation"),
+      ],
+  )
+  async def test_close_preserves_metadata(
+      self, test_id, use_progressive_sse, metadata_type
+  ):
+    """close() should carry metadata into the aggregated response."""
+    aggregator = streaming_utils.StreamingResponseAggregator()
+
+    metadata = None
+    response1 = None
+    response2 = None
+
+    if metadata_type == "grounding":
+      metadata = types.GroundingMetadata(
+          grounding_chunks=[
+              types.GroundingChunk(
+                  retrieved_context=types.GroundingChunkRetrievedContext(
+                      uri="https://example.com/doc1",
+                      title="Source",
+                  )
+              )
+          ],
+      )
+      response1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[types.Part(text="Hello ")]),
+                  grounding_metadata=metadata,
+              )
+          ]
+      )
+      response2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[types.Part(text="World!")]),
+                  finish_reason=types.FinishReason.STOP,
+                  grounding_metadata=metadata,
+              )
+          ]
+      )
+    elif metadata_type == "citation":
+      metadata = types.CitationMetadata(
+          citations=[
+              types.Citation(
+                  start_index=0,
+                  end_index=10,
+                  uri="https://example.com/source",
+                  title="Source",
+              )
+          ]
+      )
+      response1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[types.Part(text="Cited text")]),
+              )
+          ]
+      )
+      response2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[]),
+                  finish_reason=types.FinishReason.STOP,
+                  citation_metadata=metadata,
+              )
+          ]
+      )
+
+    async def run_test():
+      async for _ in aggregator.process_response(response1):
+        pass
+      async for _ in aggregator.process_response(response2):
+        pass
+
+      closed_response = aggregator.close()
+      assert closed_response is not None
+      if use_progressive_sse:
+        assert closed_response.partial is False
+
+      if metadata_type == "grounding":
+        assert closed_response.grounding_metadata is not None
+        assert len(closed_response.grounding_metadata.grounding_chunks) == 1
+      elif metadata_type == "citation":
+        assert closed_response.citation_metadata is not None
+        assert len(closed_response.citation_metadata.citations) == 1
+
+    if use_progressive_sse:
+      with temporary_feature_override(
+          FeatureName.PROGRESSIVE_SSE_STREAMING, True
+      ):
+        await run_test()
+    else:
+      await run_test()
+
+  @pytest.mark.asyncio
+  @pytest.mark.parametrize("use_progressive_sse", [False, True])
+  async def test_close_propagates_model_version(self, use_progressive_sse):
+    """close() should carry model_version into the aggregated response."""
+    aggregator = streaming_utils.StreamingResponseAggregator()
+    response1 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="Hello ")]),
+            )
+        ],
+        model_version="gemini-test-1.0",
+    )
+    response2 = types.GenerateContentResponse(
+        candidates=[
+            types.Candidate(
+                content=types.Content(parts=[types.Part(text="World!")]),
+                finish_reason=types.FinishReason.STOP,
+            )
+        ],
+        model_version="gemini-test-1.0",
+    )
+
+    async def run_test():
+      async for _ in aggregator.process_response(response1):
+        pass
+      async for _ in aggregator.process_response(response2):
+        pass
+
+      closed_response = aggregator.close()
+      assert closed_response is not None
+      assert closed_response.model_version == "gemini-test-1.0"
+
+    if use_progressive_sse:
+      with temporary_feature_override(
+          FeatureName.PROGRESSIVE_SSE_STREAMING, True
+      ):
+        await run_test()
+    else:
+      await run_test()
+
+  @pytest.mark.asyncio
+  async def test_non_progressive_merged_yield_propagates_model_version(self):
+    """The mid-stream merged text event should carry model_version forward.
+
+    In non-progressive mode, when a new non-text response arrives after buffered
+    text, the aggregator yields a synthesized merged-text LlmResponse before
+    yielding the current partial. That merged event must preserve fields from
+    the source response (model_version, grounding_metadata, citation_metadata,
+    finish_reason).
+    """
+    # PROGRESSIVE_SSE_STREAMING defaults to on; explicitly disable it to
+    # exercise the non-progressive merged-yield code path under test.
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, False
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+      # First: buffer some text.
+      response1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[types.Part(text="Hello World!")]
+                  ),
+              )
+          ],
+          model_version="gemini-test-2.0",
+      )
+      # Second: a response without text triggers the merged yield path.
+      response2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[]),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ],
+          model_version="gemini-test-2.0",
+      )
+
+      results = []
+      async for r in aggregator.process_response(response1):
+        results.append(r)
+      async for r in aggregator.process_response(response2):
+        results.append(r)
+
+      # The synthesized merged-text event should carry model_version.
+      merged_events = [
+          r
+          for r in results
+          if r.content
+          and r.content.parts
+          and r.content.parts[0].text == "Hello World!"
+          and not r.partial
+      ]
+      assert merged_events, "expected a merged non-partial text event"
+      assert merged_events[0].model_version == "gemini-test-2.0"
+
+
+class TestFunctionCallIdGeneration:
+  """Tests for function call ID generation in streaming mode."""
+
+  @pytest.mark.asyncio
+  async def test_non_streaming_fc_generates_id_when_empty(self):
+    """Non-streaming function call should get an adk-* ID if LLM didn't provide one."""
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, True
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+
+      response = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name="my_tool",
+                                  args={"x": 1},
+                                  id=None,  # No ID from LLM
+                              )
+                          )
+                      ]
+                  ),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ]
+      )
+
+      async for _ in aggregator.process_response(response):
+        pass
+
+      closed_response = aggregator.close()
+      assert closed_response is not None
+      fc = closed_response.content.parts[0].function_call
+      assert fc.id is not None
+      assert fc.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
+
+  @pytest.mark.asyncio
+  async def test_non_streaming_fc_preserves_llm_assigned_id(self):
+    """Non-streaming function call should preserve ID if LLM provided one."""
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, True
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+
+      response = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name="my_tool",
+                                  args={"x": 1},
+                                  id="llm-assigned-id",
+                              )
+                          )
+                      ]
+                  ),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ]
+      )
+
+      async for _ in aggregator.process_response(response):
+        pass
+
+      closed_response = aggregator.close()
+      assert closed_response is not None
+      fc = closed_response.content.parts[0].function_call
+      assert fc.id == "llm-assigned-id"
+
+  @pytest.mark.asyncio
+  async def test_streaming_fc_generates_consistent_id_across_chunks(self):
+    """Streaming function call should have the same ID in partial and final responses."""
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, True
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+
+      # First chunk: function call starts
+      response1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name="my_tool",
+                                  id=None,
+                                  partial_args=[
+                                      types.PartialArg(
+                                          json_path="$.x",
+                                          string_value="hello",
+                                      )
+                                  ],
+                                  will_continue=True,
+                              )
+                          )
+                      ]
+                  )
+              )
+          ]
+      )
+
+      # Second chunk: function call continues
+      response2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name=None,
+                                  id=None,
+                                  partial_args=[
+                                      types.PartialArg(
+                                          json_path="$.x",
+                                          string_value=" world",
+                                      )
+                                  ],
+                                  will_continue=False,  # Complete
+                              )
+                          )
+                      ]
+                  ),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ]
+      )
+
+      partial_results = []
+      async for r in aggregator.process_response(response1):
+        partial_results.append(r)
+      async for r in aggregator.process_response(response2):
+        partial_results.append(r)
+
+      closed_response = aggregator.close()
+      assert closed_response is not None
+      final_fc = closed_response.content.parts[0].function_call
+      assert final_fc.id is not None
+      assert final_fc.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
+      assert final_fc.args == {"x": "hello world"}
+
+      # Verify partial and final events share the same ID
+      partial_fc = partial_results[0].content.parts[0].function_call
+      assert (
+          partial_fc.id == final_fc.id
+      ), f"Partial FC ID ({partial_fc.id!r}) != Final FC ID ({final_fc.id!r})"
+
+  @pytest.mark.asyncio
+  async def test_multiple_streaming_fcs_get_different_ids(self):
+    """Multiple function calls arriving in separate chunks should get different IDs."""
+    with temporary_feature_override(
+        FeatureName.PROGRESSIVE_SSE_STREAMING, True
+    ):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+
+      # First FC
+      response1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name="tool_a",
+                                  id=None,
+                                  partial_args=[
+                                      types.PartialArg(
+                                          json_path="$.a", string_value="val_a"
+                                      )
+                                  ],
+                                  will_continue=False,
+                              )
+                          )
+                      ]
+                  )
+              )
+          ]
+      )
+
+      # Second FC
+      response2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(
+                      parts=[
+                          types.Part(
+                              function_call=types.FunctionCall(
+                                  name="tool_b",
+                                  id=None,
+                                  partial_args=[
+                                      types.PartialArg(
+                                          json_path="$.b", string_value="val_b"
+                                      )
+                                  ],
+                                  will_continue=False,
+                              )
+                          )
+                      ]
+                  ),
+                  finish_reason=types.FinishReason.STOP,
+              )
+          ]
+      )
+
+      async for _ in aggregator.process_response(response1):
+        pass
+      async for _ in aggregator.process_response(response2):
+        pass
+
+      closed_response = aggregator.close()
+      assert closed_response is not None
+      assert len(closed_response.content.parts) == 2
+
+      fc_a = closed_response.content.parts[0].function_call
+      fc_b = closed_response.content.parts[1].function_call
+
+      assert fc_a.id is not None
+      assert fc_b.id is not None
+      assert fc_a.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
+      assert fc_b.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
+      assert fc_a.id != fc_b.id  # Different IDs for different FCs

tests/unittests/utils/test_variant_utils.py

@@ -0,0 +1,43 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for variant_utils."""
+
+import warnings
+
+from google.adk.utils import variant_utils
+from google.adk.utils.variant_utils import GoogleLLMVariant
+
+
+def test_get_google_llm_variant_enterprise(monkeypatch):
+  monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true')
+  assert variant_utils.get_google_llm_variant() == GoogleLLMVariant.VERTEX_AI
+
+
+def test_get_google_llm_variant_vertexai_fallback(monkeypatch):
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter('always')
+    result = variant_utils.get_google_llm_variant()
+    assert result == GoogleLLMVariant.VERTEX_AI
+    assert len(w) == 1
+    assert issubclass(w[-1].category, DeprecationWarning)
+    assert 'GOOGLE_GENAI_USE_VERTEXAI is deprecated' in str(w[-1].message)
+
+
+def test_get_google_llm_variant_default(monkeypatch):
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.delenv('GOOGLE_GENAI_USE_VERTEXAI', raising=False)
+  assert variant_utils.get_google_llm_variant() == GoogleLLMVariant.GEMINI_API

tests/unittests/utils/test_vertex_ai_utils.py

@@ -0,0 +1,115 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for vertex_utils."""
+
+from unittest import mock
+import warnings
+
+from google.adk.utils import vertex_ai_utils
+import pytest
+
+
+def test_get_express_mode_api_key_value_error():
+  with pytest.raises(ValueError) as excinfo:
+    vertex_ai_utils.get_express_mode_api_key(
+        project='test-project', location=None, express_mode_api_key='key'
+    )
+  assert (
+      'Cannot specify project or location and express_mode_api_key. Either use'
+      ' project and location, or just the express_mode_api_key.'
+      in str(excinfo.value)
+  )
+  with pytest.raises(ValueError) as excinfo:
+    vertex_ai_utils.get_express_mode_api_key(
+        project=None, location='test-location', express_mode_api_key='key'
+    )
+  assert (
+      'Cannot specify project or location and express_mode_api_key. Either use'
+      ' project and location, or just the express_mode_api_key.'
+      in str(excinfo.value)
+  )
+  with pytest.raises(ValueError) as excinfo:
+    vertex_ai_utils.get_express_mode_api_key(
+        project='test-project',
+        location='test-location',
+        express_mode_api_key='key',
+    )
+  assert (
+      'Cannot specify project or location and express_mode_api_key. Either use'
+      ' project and location, or just the express_mode_api_key.'
+      in str(excinfo.value)
+  )
+
+
+@pytest.mark.parametrize(
+    (
+        'use_vertexai_env',
+        'google_api_key_env',
+        'express_mode_api_key',
+        'expected',
+    ),
+    [
+        ('true', None, 'express_key', 'express_key'),
+        ('1', 'google_key', 'express_key', 'express_key'),
+        ('true', 'google_key', None, 'google_key'),
+        ('1', None, None, None),
+        ('false', 'google_key', 'express_key', None),
+        ('0', 'google_key', None, None),
+        (None, 'google_key', 'express_key', None),
+    ],
+)
+def test_get_express_mode_api_key(
+    use_vertexai_env,
+    google_api_key_env,
+    express_mode_api_key,
+    expected,
+):
+  env_vars = {}
+  if use_vertexai_env:
+    env_vars['GOOGLE_GENAI_USE_ENTERPRISE'] = use_vertexai_env
+  if google_api_key_env:
+    env_vars['GOOGLE_API_KEY'] = google_api_key_env
+  with mock.patch.dict('os.environ', env_vars, clear=True):
+    assert (
+        vertex_ai_utils.get_express_mode_api_key(
+            project=None,
+            location=None,
+            express_mode_api_key=express_mode_api_key,
+        )
+        == expected
+    )
+
+
+def test_get_express_mode_api_key_enterprise(monkeypatch):
+  monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true')
+  monkeypatch.setenv('GOOGLE_API_KEY', 'google_key')
+  assert (
+      vertex_ai_utils.get_express_mode_api_key(None, None, None) == 'google_key'
+  )
+
+
+def test_get_express_mode_api_key_vertexai_fallback_warning(monkeypatch):
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+  monkeypatch.setenv('GOOGLE_API_KEY', 'google_key')
+  # Should trigger a deprecation warning and return the key
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter('always')
+    result = vertex_ai_utils.get_express_mode_api_key(None, None, None)
+
+    assert result == 'google_key'
+    assert len(w) == 1
+    assert issubclass(w[-1].category, DeprecationWarning)
+    assert 'GOOGLE_GENAI_USE_VERTEXAI is deprecated' in str(w[-1].message)

tests/unittests/utils/test_yaml_utils.py

@@ -0,0 +1,154 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for YAML utility functions."""
+
+from pathlib import Path
+from typing import Optional
+
+from google.adk.utils.yaml_utils import dump_pydantic_to_yaml
+from google.genai import types
+from pydantic import BaseModel
+
+
+class SimpleModel(BaseModel):
+  """Simple test model."""
+
+  name: str
+  age: int
+  active: bool
+  finish_reason: Optional[types.FinishReason] = None
+  multiline_text: Optional[str] = None
+  items: Optional[list[str]] = None
+
+
+def test_yaml_file_generation(tmp_path: Path):
+  """Test that YAML file is correctly generated."""
+  model = SimpleModel(
+      name="Alice",
+      age=30,
+      active=True,
+      finish_reason=types.FinishReason.STOP,
+  )
+  yaml_file = tmp_path / "test.yaml"
+
+  dump_pydantic_to_yaml(model, yaml_file)
+
+  assert yaml_file.read_text(encoding="utf-8") == """\
+active: true
+age: 30
+finish_reason: STOP
+name: Alice
+"""
+
+
+def test_multiline_string_pipe_style(tmp_path: Path):
+  """Test that multiline strings use | style."""
+  multiline_text = """\
+This is a long description
+that spans multiple lines
+and should be formatted with pipe style"""
+  model = SimpleModel(
+      name="Test",
+      age=25,
+      active=False,
+      multiline_text=multiline_text,
+  )
+  yaml_file = tmp_path / "test.yaml"
+
+  dump_pydantic_to_yaml(model, yaml_file)
+
+  assert yaml_file.read_text(encoding="utf-8") == """\
+active: false
+age: 25
+multiline_text: |-
+  This is a long description
+  that spans multiple lines
+  and should be formatted with pipe style
+name: Test
+"""
+
+
+def test_list_indentation(tmp_path: Path):
+  """Test that lists in mappings are properly indented."""
+  model = SimpleModel(
+      name="Test",
+      age=25,
+      active=True,
+      items=["item1", "item2", "item3"],
+  )
+  yaml_file = tmp_path / "test.yaml"
+
+  dump_pydantic_to_yaml(model, yaml_file)
+
+  expected = """\
+active: true
+age: 25
+items:
+  - item1
+  - item2
+  - item3
+name: Test
+"""
+  assert yaml_file.read_text(encoding="utf-8") == expected
+
+
+def test_empty_list_formatting(tmp_path: Path):
+  """Test that empty lists are formatted properly."""
+  model = SimpleModel(
+      name="Test",
+      age=25,
+      active=True,
+      items=[],
+  )
+  yaml_file = tmp_path / "test.yaml"
+
+  dump_pydantic_to_yaml(model, yaml_file)
+
+  expected = """\
+active: true
+age: 25
+items: []
+name: Test
+"""
+  assert yaml_file.read_text(encoding="utf-8") == expected
+
+
+def test_non_ascii_character_preservation(tmp_path: Path):
+  """Test that non-ASCII characters are preserved in YAML output."""
+  model = SimpleModel(
+      name="你好世界",  # Chinese
+      age=30,
+      active=True,
+      multiline_text="🌍 Hello World 🌏\nこんにちは世界\nHola Mundo 🌎",
+      items=["Château", "naïve", "café", "🎉"],
+  )
+  yaml_file = tmp_path / "test.yaml"
+
+  dump_pydantic_to_yaml(model, yaml_file)
+
+  assert yaml_file.read_text(encoding="utf-8") == """\
+active: true
+age: 30
+items:
+  - Château
+  - naïve
+  - café
+  - 🎉
+multiline_text: |-
+  🌍 Hello World 🌏
+  こんにちは世界
+  Hola Mundo 🌎
+name: 你好世界
+"""