# 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 GeminiContextCacheManager."""

import time
from unittest.mock import AsyncMock
from unittest.mock import MagicMock
from unittest.mock import patch

from google.adk.agents.context_cache_config import ContextCacheConfig
from google.adk.models.cache_metadata import CacheMetadata
from google.adk.models.gemini_context_cache_manager import GeminiContextCacheManager
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.genai import Client
from google.genai import types


class TestGeminiContextCacheManager:
  """Test suite for GeminiContextCacheManager."""

  def setup_method(self):
    """Set up test fixtures."""
    mock_client = AsyncMock(spec=Client)
    self.manager = GeminiContextCacheManager(mock_client)
    self.cache_config = ContextCacheConfig(
        cache_intervals=10,
        ttl_seconds=1800,
        min_tokens=0,  # Allow caching for tests
    )

  def create_llm_request(self, cache_metadata=None, contents_count=3):
    """Helper to create test LlmRequest."""
    contents = []
    for i in range(contents_count):
      contents.append(
          types.Content(
              role="user", parts=[types.Part(text=f"Test message {i}")]
          )
      )

    # Create tools for testing fingerprinting
    tools = [
        types.Tool(
            function_declarations=[
                types.FunctionDeclaration(
                    name="test_tool",
                    description="A test tool",
                    parameters=types.Schema(
                        type=types.Type.OBJECT,
                        properties={
                            "param": types.Schema(type=types.Type.STRING)
                        },
                    ),
                )
            ]
        )
    ]

    tool_config = types.ToolConfig(
        function_calling_config=types.FunctionCallingConfig(mode="AUTO")
    )

    return LlmRequest(
        model="gemini-2.5-flash",
        contents=contents,
        config=types.GenerateContentConfig(
            system_instruction="Test instruction",
            tools=tools,
            tool_config=tool_config,
        ),
        cache_config=self.cache_config,
        cache_metadata=cache_metadata,
    )

  def create_cache_metadata(
      self, invocations_used=0, expired=False, contents_count=3
  ):
    """Helper to create test CacheMetadata."""
    current_time = time.time()
    expire_time = current_time - 300 if expired else current_time + 1800

    return CacheMetadata(
        cache_name="projects/test/locations/us-central1/cachedContents/test123",
        expire_time=expire_time,
        fingerprint="test_fingerprint",
        invocations_used=invocations_used,
        contents_count=contents_count,
        created_at=current_time - 600,
    )

  def test_init(self):
    """Test manager initialization."""
    mock_client = MagicMock(spec=Client)
    manager = GeminiContextCacheManager(mock_client)
    assert manager is not None
    assert manager.genai_client == mock_client

  async def test_handle_context_caching_no_existing_cache(self):
    """Test handling context caching with no existing cache returns fingerprint-only metadata."""
    llm_request = self.create_llm_request(contents_count=5)

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      result = await self.manager.handle_context_caching(llm_request)

    assert result is not None
    # Should return fingerprint-only metadata (no active cache)
    assert result.cache_name is None
    assert result.expire_time is None
    assert result.invocations_used is None
    assert result.created_at is None
    assert result.fingerprint == "test_fp"
    assert result.contents_count == 0

    # No cache should be created
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_handle_context_caching_valid_existing_cache(self):
    """Test handling context caching with valid existing cache."""

    # Create request with existing valid cache
    existing_cache = self.create_cache_metadata(invocations_used=5)
    llm_request = self.create_llm_request(cache_metadata=existing_cache)

    with patch.object(self.manager, "_is_cache_valid", return_value=True):
      result = await self.manager.handle_context_caching(llm_request)

    assert result is not None
    # Verify that existing cache metadata is preserved (copied)
    assert result.cache_name == existing_cache.cache_name
    assert (
        result.invocations_used == existing_cache.invocations_used
    )  # Should preserve original invocations_used
    assert (
        result.expire_time == existing_cache.expire_time
    )  # Should preserve original expire_time
    assert (
        result.fingerprint == existing_cache.fingerprint
    )  # Should preserve original fingerprint
    assert (
        result.created_at == existing_cache.created_at
    )  # Should preserve original created_at

    # Verify it's a copy, not the same object
    assert result is not existing_cache

    # Should not create new cache
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_handle_context_caching_invalid_cache_fingerprint_match(self):
    """Test invalid cache with matching fingerprint creates new cache."""
    # Setup mocks
    mock_cached_content = AsyncMock()
    mock_cached_content.name = (
        "projects/test/locations/us-central1/cachedContents/new456"
    )
    self.manager.genai_client.aio.caches.create = AsyncMock(
        return_value=mock_cached_content
    )

    # Create request with invalid existing cache
    existing_cache = self.create_cache_metadata(
        invocations_used=15
    )  # Exceeds cache_intervals
    llm_request = self.create_llm_request(cache_metadata=existing_cache)
    llm_request.cacheable_contents_token_count = (
        5000  # Above Gemini's 4096 minimum for cache creation
    )

    with (
        patch.object(self.manager, "_is_cache_valid", return_value=False),
        patch.object(self.manager, "cleanup_cache") as mock_cleanup,
        patch.object(
            self.manager,
            "_generate_cache_fingerprint",
            return_value="test_fingerprint",  # Match old fingerprint
        ),
    ):

      result = await self.manager.handle_context_caching(llm_request)

    assert result is not None
    # Should create new cache when fingerprints match
    assert (
        result.cache_name
        == "projects/test/locations/us-central1/cachedContents/new456"
    )
    mock_cleanup.assert_called_once_with(existing_cache.cache_name)
    self.manager.genai_client.aio.caches.create.assert_called_once()

  async def test_create_cache_gates_on_prefix_not_full_prompt(self):
    """Cache creation is gated on the cacheable prefix, not the full prompt.

    Regression test for https://github.com/google/adk-python/issues/5847.

    On a long conversation the previous-prompt token count
    (``cacheable_contents_token_count``) can be well above Gemini's 4096-token
    minimum while the cached prefix ``contents[:cache_contents_count]`` is far
    below it. Creating a cache in that case makes ``caches.create`` fail with a
    400 INVALID_ARGUMENT. The manager must skip cache creation instead.
    """
    self.manager.genai_client.aio.caches.create = AsyncMock()

    # A tiny cacheable prefix followed by a huge trailing user turn.
    contents = [
        types.Content(role="user", parts=[types.Part(text="Short prefix.")]),
        types.Content(role="user", parts=[types.Part(text="word " * 100_000)]),
    ]
    llm_request = LlmRequest(
        model="gemini-2.5-flash",
        contents=contents,
        config=types.GenerateContentConfig(
            system_instruction="You are a helpful assistant.",
        ),
        cache_config=self.cache_config,
    )
    # Full previous prompt is large (clears the old, buggy gate)...
    llm_request.cacheable_contents_token_count = 75000

    # ...but only the tiny first content is cacheable.
    result = await self.manager._create_new_cache_with_contents(
        llm_request, cache_contents_count=1
    )

    assert result is None
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_handle_context_caching_invalid_cache_fingerprint_mismatch(
      self,
  ):
    """Test invalid cache with mismatched fingerprint returns fingerprint-only metadata."""
    # Create request with invalid existing cache
    existing_cache = self.create_cache_metadata(
        invocations_used=15, contents_count=3
    )  # Exceeds cache_intervals
    llm_request = self.create_llm_request(
        cache_metadata=existing_cache, contents_count=5
    )

    with (
        patch.object(self.manager, "_is_cache_valid", return_value=False),
        patch.object(self.manager, "cleanup_cache") as mock_cleanup,
        patch.object(
            self.manager,
            "_generate_cache_fingerprint",
            side_effect=["old_fp", "new_fp"],  # Different fingerprints
        ),
    ):

      result = await self.manager.handle_context_caching(llm_request)

    assert result is not None
    # Should return fingerprint-only metadata
    assert result.cache_name is None
    assert result.expire_time is None
    assert result.invocations_used is None
    assert result.created_at is None
    assert result.fingerprint == "new_fp"
    assert result.contents_count == 0
    mock_cleanup.assert_called_once_with(existing_cache.cache_name)
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_is_cache_valid_fingerprint_mismatch(self):
    """Test cache validation with fingerprint mismatch."""
    cache_metadata = self.create_cache_metadata()
    llm_request = self.create_llm_request(cache_metadata=cache_metadata)

    with patch.object(
        self.manager,
        "_generate_cache_fingerprint",
        return_value="different_fingerprint",
    ):
      result = await self.manager._is_cache_valid(llm_request)

    assert result is False

  async def test_is_cache_valid_expired_cache(self):
    """Test cache validation with expired cache."""
    cache_metadata = self.create_cache_metadata(expired=True)
    llm_request = self.create_llm_request(cache_metadata=cache_metadata)

    with patch.object(
        self.manager,
        "_generate_cache_fingerprint",
        return_value="test_fingerprint",
    ):
      result = await self.manager._is_cache_valid(llm_request)

    assert result is False

  async def test_is_cache_valid_fingerprint_only_metadata(self):
    """Test cache validation with fingerprint-only metadata (no active cache)."""
    # Create fingerprint-only metadata (cache_name is None)
    cache_metadata = CacheMetadata(
        fingerprint="test_fingerprint",
        contents_count=5,
    )
    llm_request = self.create_llm_request(cache_metadata=cache_metadata)

    result = await self.manager._is_cache_valid(llm_request)

    assert (
        result is False
    )  # Fingerprint-only metadata is not a valid active cache

  async def test_is_cache_valid_cache_intervals_exceeded(self):
    """Test cache validation with max invocations exceeded."""
    cache_metadata = self.create_cache_metadata(
        invocations_used=15
    )  # Exceeds cache_intervals=10
    llm_request = self.create_llm_request(cache_metadata=cache_metadata)

    with patch.object(
        self.manager,
        "_generate_cache_fingerprint",
        return_value="test_fingerprint",
    ):
      result = await self.manager._is_cache_valid(llm_request)

    assert result is False

  async def test_is_cache_valid_all_checks_pass(self):
    """Test cache validation when all checks pass."""
    cache_metadata = self.create_cache_metadata(
        invocations_used=5
    )  # Within cache_intervals=10
    llm_request = self.create_llm_request(cache_metadata=cache_metadata)

    with patch.object(
        self.manager,
        "_generate_cache_fingerprint",
        return_value="test_fingerprint",
    ):
      result = await self.manager._is_cache_valid(llm_request)

    assert result is True

  async def test_cleanup_cache(self):
    """Test cache cleanup functionality."""
    cache_name = "projects/test/locations/us-central1/cachedContents/test123"

    await self.manager.cleanup_cache(cache_name)

    self.manager.genai_client.aio.caches.delete.assert_called_once_with(
        name=cache_name
    )

  def test_generate_cache_fingerprint(self):
    """Test cache fingerprint generation includes tools and tool_config."""
    llm_request = self.create_llm_request()
    cache_contents_count = 2  # Cache all but last content

    fingerprint1 = self.manager._generate_cache_fingerprint(
        llm_request, cache_contents_count
    )
    fingerprint2 = self.manager._generate_cache_fingerprint(
        llm_request, cache_contents_count
    )

    # Same request should generate same fingerprint
    assert fingerprint1 == fingerprint2
    assert isinstance(fingerprint1, str)
    assert len(fingerprint1) > 0

    # Test that tool_config and tools are included in fingerprint
    # Create request without tools/tool_config
    llm_request_no_tools = LlmRequest(
        model="gemini-2.5-flash",
        contents=[types.Content(role="user", parts=[types.Part(text="Test")])],
        config=types.GenerateContentConfig(
            system_instruction="Test instruction"
        ),
        cache_config=self.cache_config,
    )

    fingerprint_no_tools = self.manager._generate_cache_fingerprint(
        llm_request_no_tools, cache_contents_count
    )

    # Should be different from request with tools
    assert fingerprint1 != fingerprint_no_tools

  def test_generate_cache_fingerprint_different_requests(self):
    """Test that different requests generate different fingerprints."""
    llm_request1 = self.create_llm_request()

    llm_request2 = LlmRequest(
        model="gemini-2.5-flash",
        contents=[
            types.Content(
                role="user", parts=[types.Part(text="Different message")]
            )
        ],
        config=types.GenerateContentConfig(
            system_instruction="Different instruction"
        ),
        cache_config=self.cache_config,
    )

    cache_contents_count = 2
    fingerprint1 = self.manager._generate_cache_fingerprint(
        llm_request1, cache_contents_count
    )
    fingerprint2 = self.manager._generate_cache_fingerprint(
        llm_request2, cache_contents_count
    )

    assert fingerprint1 != fingerprint2

  def test_generate_cache_fingerprint_tool_config_variations(self):
    """Test that different tool configs generate different fingerprints."""
    # Request with AUTO mode
    llm_request_auto = self.create_llm_request()

    # Request with NONE mode
    tool_config_none = types.ToolConfig(
        function_calling_config=types.FunctionCallingConfig(mode="NONE")
    )

    llm_request_none = LlmRequest(
        model="gemini-2.5-flash",
        contents=[types.Content(role="user", parts=[types.Part(text="Test")])],
        config=types.GenerateContentConfig(
            system_instruction="Test instruction",
            tools=llm_request_auto.config.tools,
            tool_config=tool_config_none,
        ),
        cache_config=self.cache_config,
    )

    cache_contents_count = 2
    fingerprint_auto = self.manager._generate_cache_fingerprint(
        llm_request_auto, cache_contents_count
    )
    fingerprint_none = self.manager._generate_cache_fingerprint(
        llm_request_none, cache_contents_count
    )

    assert fingerprint_auto != fingerprint_none

  async def test_populate_cache_metadata_in_response_no_invocations_increment(
      self,
  ):
    """Test that populate_cache_metadata_in_response doesn't increment invocations_used."""
    # Create mock response with usage metadata
    usage_metadata = MagicMock()
    usage_metadata.cached_content_token_count = 800
    usage_metadata.prompt_token_count = 1000

    llm_response = MagicMock(spec=LlmResponse)
    llm_response.usage_metadata = usage_metadata

    cache_metadata = self.create_cache_metadata(invocations_used=3)

    self.manager.populate_cache_metadata_in_response(
        llm_response, cache_metadata
    )

    # Verify response metadata preserves the original invocations_used (no increment)
    updated_metadata = llm_response.cache_metadata
    assert (
        updated_metadata.invocations_used == 3
    )  # Should preserve original value
    assert updated_metadata.cache_name == cache_metadata.cache_name
    assert updated_metadata.fingerprint == cache_metadata.fingerprint
    assert updated_metadata.expire_time == cache_metadata.expire_time
    assert updated_metadata.created_at == cache_metadata.created_at

  async def test_populate_cache_metadata_no_usage_metadata(self):
    """Test populating cache metadata when no usage metadata."""
    llm_response = MagicMock(spec=LlmResponse)
    llm_response.usage_metadata = None

    cache_metadata = self.create_cache_metadata(invocations_used=3)

    self.manager.populate_cache_metadata_in_response(
        llm_response, cache_metadata
    )

    # Should still create metadata even without usage info
    updated_metadata = llm_response.cache_metadata
    assert (
        updated_metadata.invocations_used == 3
    )  # Should preserve original value
    assert updated_metadata.cache_name == cache_metadata.cache_name

  async def test_create_new_cache_with_proper_ttl(self):
    """Test that new cache is created with proper TTL."""
    mock_cached_content = AsyncMock()
    mock_cached_content.name = (
        "projects/test/locations/us-central1/cachedContents/test123"
    )
    self.manager.genai_client.aio.caches.create = AsyncMock(
        return_value=mock_cached_content
    )

    llm_request = self.create_llm_request()

    cache_contents_count = max(0, len(llm_request.contents) - 1)

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      await self.manager._create_gemini_cache(llm_request, cache_contents_count)

    # Verify cache creation call includes TTL
    create_call = self.manager.genai_client.aio.caches.create.call_args
    assert create_call is not None
    cache_config = create_call[1]["config"]
    assert cache_config.ttl == "1800s"  # From cache_config

  def test_all_but_last_content_caching(self):
    """Test that cache content counting works correctly."""
    # Test with multiple contents
    llm_request_multi = self.create_llm_request(contents_count=5)

    # Test cache contents count calculation
    cache_contents_count = max(0, len(llm_request_multi.contents) - 1)

    assert cache_contents_count == 4  # 5 contents, so cache 4 contents

    # Test with single content
    llm_request_single = self.create_llm_request(contents_count=1)
    single_cache_contents_count = max(0, len(llm_request_single.contents) - 1)

    assert single_cache_contents_count == 0  # Single content, cache 0 contents

  def test_edge_cases(self):
    """Test various edge cases."""
    # Test with None cache_config
    llm_request_no_config = LlmRequest(
        model="gemini-2.5-flash",
        contents=[types.Content(role="user", parts=[types.Part(text="Test")])],
        config=types.GenerateContentConfig(system_instruction="Test"),
        cache_config=None,
    )

    # Should handle gracefully
    cache_contents_count = 2
    fingerprint = self.manager._generate_cache_fingerprint(
        llm_request_no_config, cache_contents_count
    )
    assert isinstance(fingerprint, str)

    # Test with empty contents
    llm_request_empty = LlmRequest(
        model="gemini-2.5-flash",
        contents=[],
        config=types.GenerateContentConfig(system_instruction="Test"),
        cache_config=self.cache_config,
    )

    empty_cache_contents_count = 0
    fingerprint = self.manager._generate_cache_fingerprint(
        llm_request_empty, empty_cache_contents_count
    )
    assert isinstance(fingerprint, str)

  def test_parameter_types_enforcement(self):
    """Test that method calls with correct parameter types work properly."""
    # Create proper objects
    usage_metadata = MagicMock()
    usage_metadata.cached_content_token_count = 500
    usage_metadata.prompt_token_count = 1000

    llm_response = MagicMock(spec=LlmResponse)
    llm_response.usage_metadata = usage_metadata

    cache_metadata = self.create_cache_metadata(invocations_used=3)

    # This should work fine (correct types and order)
    self.manager.populate_cache_metadata_in_response(
        llm_response, cache_metadata
    )
    updated_metadata = llm_response.cache_metadata
    assert updated_metadata.invocations_used == 3  # No increment in this method

    # Document expected types for integration tests
    assert isinstance(cache_metadata, CacheMetadata)
    assert hasattr(
        llm_response, "usage_metadata"
    )  # LlmResponse should have this
    assert not hasattr(
        cache_metadata, "usage_metadata"
    )  # CacheMetadata should NOT have this

  def create_llm_request_with_token_count(
      self, token_count=None, cache_metadata=None
  ):
    """Helper to create LlmRequest with cacheable_contents_token_count."""
    llm_request = self.create_llm_request(cache_metadata=cache_metadata)
    llm_request.cacheable_contents_token_count = token_count
    return llm_request

  async def test_cache_creation_with_sufficient_token_count(self):
    """Test that fingerprint-only metadata is returned even with sufficient tokens."""
    # With new prefix matching logic, no cache is created without existing metadata
    # Create request with sufficient token count
    llm_request = self.create_llm_request_with_token_count(token_count=2048)

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      result = await self.manager.handle_context_caching(llm_request)

    # Should return fingerprint-only metadata (no cache creation)
    assert result is not None
    assert result.cache_name is None  # Fingerprint-only state
    assert result.fingerprint == "test_fp"
    assert result.contents_count == 0
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_cache_creation_with_insufficient_token_count(self):
    """Test that fingerprint-only metadata is returned even with insufficient tokens."""
    # Set higher minimum token requirement
    self.manager.cache_config = ContextCacheConfig(
        cache_intervals=10,
        ttl_seconds=1800,
        min_tokens=2048,
    )

    # Create request with insufficient token count
    llm_request = self.create_llm_request_with_token_count(token_count=1024)
    llm_request.cache_config = self.manager.cache_config

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      result = await self.manager.handle_context_caching(llm_request)

    # Should return fingerprint-only metadata
    assert result is not None
    assert result.cache_name is None
    assert result.fingerprint == "test_fp"
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_cache_creation_without_token_count(self):
    """Test that fingerprint-only metadata is returned even without token count."""
    # Create request without token count (initial request)
    llm_request = self.create_llm_request_with_token_count(token_count=None)

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      result = await self.manager.handle_context_caching(llm_request)

    # Should return fingerprint-only metadata
    assert result is not None
    assert result.cache_name is None
    assert result.fingerprint == "test_fp"
    self.manager.genai_client.aio.caches.create.assert_not_called()

  async def test_fingerprint_stability_across_growing_contents_within_invocation(
      self,
  ):
    """Fingerprint over a prefix stays stable as contents grow.

    Within a single invocation, contents grow as tool calls happen:
    [user_msg] -> [user_msg, model_tool_call, tool_response].
    A fingerprint computed over contents[:1] should be the same
    regardless of how many entries follow.
    """
    user_msg = types.Content(
        role="user", parts=[types.Part(text="What is the weather?")]
    )
    model_tool_call = types.Content(
        role="model",
        parts=[
            types.Part(
                function_call=types.FunctionCall(
                    name="get_weather", args={"city": "NYC"}
                )
            )
        ],
    )
    tool_response = types.Content(
        role="user",
        parts=[
            types.Part(
                function_response=types.FunctionResponse(
                    name="get_weather", response={"temp": "72F"}
                )
            )
        ],
    )

    # First LLM call: contents = [user_msg]
    request_short = LlmRequest(
        model="gemini-2.5-flash",
        contents=[user_msg],
        config=types.GenerateContentConfig(
            system_instruction="You are a weather bot",
        ),
        cache_config=self.cache_config,
    )
    fp_short = self.manager._generate_cache_fingerprint(request_short, 1)

    # Second LLM call: contents grew to [user_msg, model, tool_resp]
    request_long = LlmRequest(
        model="gemini-2.5-flash",
        contents=[user_msg, model_tool_call, tool_response],
        config=types.GenerateContentConfig(
            system_instruction="You are a weather bot",
        ),
        cache_config=self.cache_config,
    )
    fp_long = self.manager._generate_cache_fingerprint(
        request_long, 1  # Still fingerprint over first 1 content
    )

    # Fingerprints over the same prefix must be identical
    assert fp_short == fp_long

  async def test_fingerprint_preserved_on_cache_creation_failure(self):
    """When cache creation fails, contents_count is preserved.

    When _create_new_cache_with_contents returns None (e.g., no token
    count or below Gemini's 4096 minimum), the code preserves the
    original contents_count so the fingerprint stays stable for
    subsequent calls.
    """
    # Simulate first call returning fingerprint-only metadata
    # with contents_count=3 (the original prefix size)
    first_metadata = CacheMetadata(
        fingerprint="fp_for_3",
        contents_count=3,
    )

    # Second call: contents grew to 5 entries but we carry forward
    # old metadata with contents_count=3
    llm_request = self.create_llm_request(
        cache_metadata=first_metadata, contents_count=5
    )
    llm_request.cacheable_contents_token_count = None  # No token count

    with patch.object(
        self.manager,
        "_generate_cache_fingerprint",
        side_effect=lambda _req, count: f"fp_for_{count}",
    ):
      result = await self.manager.handle_context_caching(llm_request)

    # Fix: contents_count and fingerprint are preserved from the
    # original prefix, not reset to total array length.
    assert result.cache_name is None
    assert result.contents_count == 3
    assert result.fingerprint == "fp_for_3"

  async def test_multi_turn_fingerprint_stable_when_below_token_threshold(
      self,
  ):
    """Fingerprint stays stable across turns when cache creation fails.

    Simulates 3 invocations where cache creation always fails because
    there is no token count. After the fix, contents_count is preserved
    so the fingerprint remains stable across calls.
    """
    fingerprints_seen = []
    contents_counts_seen = []
    metadata = None

    for turn in range(3):
      contents_count = 1 + turn * 2  # 1, 3, 5
      llm_request = self.create_llm_request(
          cache_metadata=metadata,
          contents_count=contents_count,
      )
      llm_request.cacheable_contents_token_count = None

      result = await self.manager.handle_context_caching(llm_request)

      assert result is not None
      assert result.cache_name is None
      fingerprints_seen.append(result.fingerprint)
      contents_counts_seen.append(result.contents_count)
      metadata = result

    # All contents in this helper are user-role messages, so there is no
    # cacheable content prefix before the final user batch.
    assert len(set(fingerprints_seen)) == 1
    assert contents_counts_seen == [0, 0, 0]

  async def test_contents_count_should_remain_stable_after_cache_creation_failure(
      self,
  ):
    """Preserved contents_count keeps fingerprint stable on failure.

    When cache creation fails, the returned metadata preserves the
    original contents_count from the prefix, not reset to the total
    number of contents. This keeps the fingerprint stable across
    LLM calls within the same invocation.
    """
    # First call: fingerprint-only metadata with contents_count=2
    first_metadata = CacheMetadata(
        fingerprint="original_fp",
        contents_count=2,
    )

    # Second call: contents grew to 5 but old metadata says 2
    llm_request = self.create_llm_request(
        cache_metadata=first_metadata, contents_count=5
    )
    llm_request.cacheable_contents_token_count = None

    # Use real fingerprint generation so the prefix fingerprint
    # matches the old metadata's fingerprint
    original_fp = self.manager._generate_cache_fingerprint(llm_request, 2)
    first_metadata = CacheMetadata(
        fingerprint=original_fp,
        contents_count=2,
    )
    llm_request.cache_metadata = first_metadata

    result = await self.manager.handle_context_caching(llm_request)

    # EXPECTED: contents_count should stay at 2 (the prefix size)
    assert result.contents_count == 2
    # EXPECTED: fingerprint should match the original
    assert result.fingerprint == original_fp

  def test_multi_tool_call_single_invocation_contents_growth(self):
    """Test _find_count_of_contents_to_cache with tool call pattern.

    Simulates realistic contents growth within a single invocation:
    user_msg -> model_tool_call -> tool_response -> model_tool_call
    -> tool_response -> final_model_response.
    """
    user_msg = types.Content(
        role="user",
        parts=[types.Part(text="Find weather and news")],
    )
    model_tool_call_1 = types.Content(
        role="model",
        parts=[
            types.Part(
                function_call=types.FunctionCall(
                    name="get_weather", args={"city": "NYC"}
                )
            )
        ],
    )
    tool_response_1 = types.Content(
        role="user",
        parts=[
            types.Part(
                function_response=types.FunctionResponse(
                    name="get_weather", response={"temp": "72F"}
                )
            )
        ],
    )
    model_tool_call_2 = types.Content(
        role="model",
        parts=[
            types.Part(
                function_call=types.FunctionCall(
                    name="get_news", args={"topic": "tech"}
                )
            )
        ],
    )
    tool_response_2 = types.Content(
        role="user",
        parts=[
            types.Part(
                function_response=types.FunctionResponse(
                    name="get_news", response={"headline": "AI advances"}
                )
            )
        ],
    )
    final_model_response = types.Content(
        role="model",
        parts=[types.Part(text="Weather is 72F, news: AI advances")],
    )

    # Stage 1: Just user message
    contents_1 = [user_msg]
    count_1 = self.manager._find_count_of_contents_to_cache(contents_1)
    assert count_1 == 0  # Only user content, nothing to cache before

    # Stage 2: After first tool call cycle
    contents_2 = [user_msg, model_tool_call_1, tool_response_1]
    count_2 = self.manager._find_count_of_contents_to_cache(contents_2)
    # Last user batch is tool_response_1 at index 2
    # model_tool_call_1 at index 1 breaks the batch
    # So cache everything before index 2 = 2 items
    assert count_2 == 2

    # Stage 3: After second tool call cycle
    contents_3 = [
        user_msg,
        model_tool_call_1,
        tool_response_1,
        model_tool_call_2,
        tool_response_2,
    ]
    count_3 = self.manager._find_count_of_contents_to_cache(contents_3)
    # Last user batch is tool_response_2 at index 4
    # model_tool_call_2 at index 3 breaks the batch
    # So cache everything before index 4 = 4 items
    assert count_3 == 4

    # Stage 4: After final model response
    contents_4 = [
        user_msg,
        model_tool_call_1,
        tool_response_1,
        model_tool_call_2,
        tool_response_2,
        final_model_response,
    ]
    count_4 = self.manager._find_count_of_contents_to_cache(contents_4)
    # Last entry is model content, no trailing user batch
    # All contents are before the (empty) last user batch
    assert count_4 == 6

  async def test_fingerprint_only_metadata_transitions_to_active_cache(
      self,
  ):
    """Happy path: fingerprint-only transitions to active cache.

    Simulates the full lifecycle across two LLM calls within the
    same invocation using real fingerprint generation:
    1. First call: no metadata -> returns fingerprint-only metadata
    2. Second call: fingerprint matches, cache created successfully
    """
    # --- First LLM call: no existing metadata ---
    llm_request_1 = self.create_llm_request(contents_count=3)

    result_1 = await self.manager.handle_context_caching(llm_request_1)

    assert result_1 is not None
    assert result_1.cache_name is None
    assert result_1.contents_count == 0

    # --- Second LLM call: carry forward fingerprint-only metadata ---
    # Contents grew but we still have same prefix
    llm_request_2 = self.create_llm_request(
        cache_metadata=result_1, contents_count=5
    )
    # contents_count is 0 (all-user conversation), so the cached prefix is the
    # system instruction + tools; use a large previous-prompt count so the
    # estimated prefix clears Gemini's 4096-token minimum.
    llm_request_2.cacheable_contents_token_count = 30000

    # Verify prefix fingerprint matches (real implementation).
    # The fingerprint-only metadata is "invalid" (no cache_name),
    # so _is_cache_valid returns False. Then the code checks if
    # the prefix fingerprint matches before attempting cache creation.
    prefix_fp = self.manager._generate_cache_fingerprint(
        llm_request_2, result_1.contents_count
    )
    assert prefix_fp == result_1.fingerprint, (
        f"Prefix fingerprint mismatch: {prefix_fp!r} != "
        f"{result_1.fingerprint!r}. "
        "This indicates the contents_count was not preserved."
    )

    # Fingerprints match - cache creation should be attempted
    mock_cached_content = AsyncMock()
    mock_cached_content.name = (
        "projects/test/locations/us-central1/cachedContents/new789"
    )
    self.manager.genai_client.aio.caches.create = AsyncMock(
        return_value=mock_cached_content
    )

    result_2 = await self.manager.handle_context_caching(llm_request_2)

    assert result_2 is not None
    assert result_2.cache_name == (
        "projects/test/locations/us-central1/cachedContents/new789"
    )
    assert result_2.contents_count == 0  # Preserved from prefix
    assert result_2.invocations_used == 1
    self.manager.genai_client.aio.caches.create.assert_called_once()

  async def test_dynamic_instruction_does_not_break_initial_cache_fingerprint(
      self,
  ):
    """Request-scoped dynamic instructions stay out of the cache prefix."""
    dynamic_instruction = types.Content(
        role="user", parts=[types.Part(text="Turn context: locale=en-US")]
    )
    user_msg = types.Content(
        role="user", parts=[types.Part(text="what time is it?")]
    )
    model_tool_call = types.Content(
        role="model",
        parts=[
            types.Part(
                function_call=types.FunctionCall(name="get_time", args={})
            )
        ],
    )
    tool_response = types.Content(
        role="user",
        parts=[
            types.Part(
                function_response=types.FunctionResponse(
                    name="get_time", response={"time": "12:00"}
                )
            )
        ],
    )

    request_1 = self.create_llm_request(contents_count=0)
    request_1.contents = [dynamic_instruction, user_msg]

    result_1 = await self.manager.handle_context_caching(request_1)

    assert result_1 is not None
    assert result_1.cache_name is None
    assert result_1.contents_count == 0

    request_2 = self.create_llm_request(
        cache_metadata=result_1, contents_count=0
    )
    request_2.contents = [
        user_msg,
        model_tool_call,
        dynamic_instruction,
        tool_response,
    ]
    # contents_count is 0, so the cached prefix is the system instruction +
    # tools; use a large previous-prompt count so the estimated prefix clears
    # Gemini's 4096-token minimum.
    request_2.cacheable_contents_token_count = 30000

    mock_cached_content = AsyncMock()
    mock_cached_content.name = (
        "projects/test/locations/us-central1/cachedContents/new789"
    )
    self.manager.genai_client.aio.caches.create = AsyncMock(
        return_value=mock_cached_content
    )

    result_2 = await self.manager.handle_context_caching(request_2)

    assert result_2 is not None
    assert result_2.cache_name == (
        "projects/test/locations/us-central1/cachedContents/new789"
    )
    assert result_2.contents_count == 0
    assert result_2.invocations_used == 1
    self.manager.genai_client.aio.caches.create.assert_called_once()

  async def test_create_http_options_passthrough(self):
    """Test that create_http_options is passed through to cache creation config."""
    mock_cached_content = AsyncMock()
    mock_cached_content.name = (
        "projects/test/locations/us-central1/cachedContents/test123"
    )
    self.manager.genai_client.aio.caches.create = AsyncMock(
        return_value=mock_cached_content
    )

    # Create config with http_options (e.g. 10s timeout)
    http_options = types.HttpOptions(timeout=10000)
    cache_config_with_timeout = ContextCacheConfig(
        cache_intervals=10,
        ttl_seconds=1800,
        min_tokens=0,
        create_http_options=http_options,
    )

    llm_request = self.create_llm_request()
    llm_request.cache_config = cache_config_with_timeout

    cache_contents_count = max(0, len(llm_request.contents) - 1)

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      await self.manager._create_gemini_cache(llm_request, cache_contents_count)

    # Verify cache creation call includes http_options
    create_call = self.manager.genai_client.aio.caches.create.call_args
    assert create_call is not None
    cache_config = create_call[1]["config"]
    assert cache_config.http_options is not None
    assert cache_config.http_options.timeout == 10000

  async def test_create_without_http_options(self):
    """Test that cache creation works without create_http_options."""
    mock_cached_content = AsyncMock()
    mock_cached_content.name = (
        "projects/test/locations/us-central1/cachedContents/test123"
    )
    self.manager.genai_client.aio.caches.create = AsyncMock(
        return_value=mock_cached_content
    )

    llm_request = self.create_llm_request()
    cache_contents_count = max(0, len(llm_request.contents) - 1)

    with patch.object(
        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
    ):
      await self.manager._create_gemini_cache(llm_request, cache_contents_count)

    # Verify cache creation call does not include http_options
    create_call = self.manager.genai_client.aio.caches.create.call_args
    assert create_call is not None
    cache_config = create_call[1]["config"]
    assert cache_config.http_options is None
