fix: preserve cache fingerprint stability on creation failure

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 905207529

Author: GWeale

src/google/adk/models/gemini_context_cache_manager.py

@@ -32,6 +32,9 @@
 
 logger = logging.getLogger("google_adk." + __name__)
 
+# Gemini API requires a minimum of 4096 tokens for cached content.
+_GEMINI_MIN_CACHE_TOKENS = 4096
+
 if TYPE_CHECKING:
   from google.genai import Client
 
@@ -119,6 +122,19 @@ async def handle_context_caching(
             )
             return cache_metadata
 
+          # Cache creation failed (e.g., below Gemini's 4096 token minimum).
+          # Preserve the original contents_count so the fingerprint stays
+          # stable for subsequent calls instead of resetting to total.
+          logger.debug(
+              "Cache creation failed, preserving prefix fingerprint "
+              "(contents_count=%d)",
+              cache_contents_count,
+          )
+          return CacheMetadata(
+              fingerprint=current_fingerprint,
+              contents_count=cache_contents_count,
+          )
+
         # Fingerprints don't match - recalculate with total contents
         logger.debug(
             "Fingerprints don't match, returning fingerprint-only metadata"
@@ -304,6 +320,15 @@ async def _create_new_cache_with_contents(
       )
       return None
 
+    # Check client-side to avoid unnecessary API round-trips.
+    if llm_request.cacheable_contents_token_count < _GEMINI_MIN_CACHE_TOKENS:
+      logger.info(
+          "Request below Gemini minimum cache size (%d < %d tokens)",
+          llm_request.cacheable_contents_token_count,
+          _GEMINI_MIN_CACHE_TOKENS,
+      )
+      return None
+
     try:
       # Create cache using Gemini API directly
       return await self._create_gemini_cache(llm_request, cache_contents_count)

tests/unittests/agents/test_gemini_context_cache_manager.py

@@ -26,7 +26,6 @@
 from google.adk.models.llm_response import LlmResponse
 from google.genai import Client
 from google.genai import types
-import pytest
 
 
 class TestGeminiContextCacheManager:
@@ -179,7 +178,7 @@ async def test_handle_context_caching_invalid_cache_fingerprint_match(self):
     )  # Exceeds cache_intervals
     llm_request = self.create_llm_request(cache_metadata=existing_cache)
     llm_request.cacheable_contents_token_count = (
-        2048  # Add token count for cache creation
+        5000  # Above Gemini's 4096 minimum for cache creation
     )
 
     with (
@@ -627,3 +626,328 @@ async def test_cache_creation_without_token_count(self):
     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.0-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.0-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
+
+    # First turn has no metadata, so uses total (1).
+    # Subsequent turns preserve contents_count=1 from the prefix.
+    # Fingerprint stays stable because contents[:1] is always the
+    # same user message.
+    assert len(set(fingerprints_seen)) == 1
+    assert contents_counts_seen == [1, 1, 1]
+
+  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 == 3
+
+    # --- 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
+    )
+    llm_request_2.cacheable_contents_token_count = 4096
+
+    # 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 == 3  # Preserved from prefix
+    assert result_2.invocations_used == 1
+    self.manager.genai_client.aio.caches.create.assert_called_once()