fix: Create context cache only when prefix matches with previous request

PiperOrigin-RevId: 817468275

tests/unittests/agents/test_gemini_context_cache_manager.py

@@ -87,7 +87,7 @@ class TestGeminiContextCacheManager:
     )
 
   def create_cache_metadata(
-      self, invocations_used=0, expired=False, cached_contents_count=3
+      self, invocations_used=0, expired=False, contents_count=3
   ):
     """Helper to create test CacheMetadata."""
     current_time = time.time()
@@ -98,7 +98,7 @@ class TestGeminiContextCacheManager:
         expire_time=expire_time,
         fingerprint="test_fingerprint",
         invocations_used=invocations_used,
-        cached_contents_count=cached_contents_count,
+        contents_count=contents_count,
         created_at=current_time - 600,
     )
 
@@ -109,45 +109,26 @@ class TestGeminiContextCacheManager:
     assert manager is not None
     assert manager.genai_client == mock_client
 
-  async def test_handle_context_caching_new_cache(self):
-    """Test handling context caching with no existing cache."""
-    # Setup mocks
-    mock_cached_content = AsyncMock()
-    mock_cached_content.name = (
-        "projects/test/locations/us-central1/cachedContents/new123"
-    )
-    self.manager.genai_client.aio.caches.create = AsyncMock(
-        return_value=mock_cached_content
-    )
-
-    llm_request = self.create_llm_request()
-    llm_request.cacheable_contents_token_count = (
-        2048  # Add token count for cache creation
-    )
-    start_time = time.time()
+  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)
 
-    end_time = time.time()
-
     assert result is not None
-    # Verify new cache metadata is created with fresh values
-    assert (
-        result.cache_name
-        == "projects/test/locations/us-central1/cachedContents/new123"
-    )
-    assert result.invocations_used == 1  # New cache starts with 1 invocation
+    # 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 == 5  # Total contents count
 
-    # Verify timestamps are recent (within test execution time)
-    assert start_time <= result.created_at <= end_time
-    assert result.expire_time > time.time()  # Should be in the future
-
-    # Verify cache creation was called
-    self.manager.genai_client.aio.caches.create.assert_called_once()
+    # 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."""
@@ -181,8 +162,8 @@ class TestGeminiContextCacheManager:
     # Should not create new cache
     self.manager.genai_client.aio.caches.create.assert_not_called()
 
-  async def test_handle_context_caching_invalid_existing_cache(self):
-    """Test handling context caching with invalid existing cache."""
+  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 = (
@@ -205,13 +186,16 @@ class TestGeminiContextCacheManager:
         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="new_fp"
+            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"
@@ -219,6 +203,41 @@ class TestGeminiContextCacheManager:
     mock_cleanup.assert_called_once_with(existing_cache.cache_name)
     self.manager.genai_client.aio.caches.create.assert_called_once()
 
+  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 == 5  # Total contents count
+    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()
@@ -247,6 +266,21 @@ class TestGeminiContextCacheManager:
 
     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(
@@ -537,16 +571,8 @@ class TestGeminiContextCacheManager:
     return llm_request
 
   async def test_cache_creation_with_sufficient_token_count(self):
-    """Test cache creation succeeds when token count meets minimum."""
-    # Setup mocks
-    mock_cached_content = AsyncMock()
-    mock_cached_content.name = (
-        "projects/test/locations/us-central1/cachedContents/token123"
-    )
-    self.manager.genai_client.aio.caches.create = AsyncMock(
-        return_value=mock_cached_content
-    )
-
+    """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)
 
@@ -555,13 +581,15 @@ class TestGeminiContextCacheManager:
     ):
       result = await self.manager.handle_context_caching(llm_request)
 
-    # Should succeed in creating cache
+    # Should return fingerprint-only metadata (no cache creation)
     assert result is not None
-    assert result.cache_name == mock_cached_content.name
-    self.manager.genai_client.aio.caches.create.assert_called_once()
+    assert result.cache_name is None  # Fingerprint-only state
+    assert result.fingerprint == "test_fp"
+    assert result.contents_count == 3
+    self.manager.genai_client.aio.caches.create.assert_not_called()
 
   async def test_cache_creation_with_insufficient_token_count(self):
-    """Test cache creation fails when token count is below minimum."""
+    """Test that fingerprint-only metadata is returned even with insufficient tokens."""
     # Set higher minimum token requirement
     self.manager.cache_config = ContextCacheConfig(
         cache_intervals=10,
@@ -573,19 +601,29 @@ class TestGeminiContextCacheManager:
     llm_request = self.create_llm_request_with_token_count(token_count=1024)
     llm_request.cache_config = self.manager.cache_config
 
-    result = await self.manager.handle_context_caching(llm_request)
+    with patch.object(
+        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
+    ):
+      result = await self.manager.handle_context_caching(llm_request)
 
-    # Should not create cache
-    assert result is None
+    # 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 cache creation is skipped when no token count is available."""
+    """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)
 
-    result = await self.manager.handle_context_caching(llm_request)
+    with patch.object(
+        self.manager, "_generate_cache_fingerprint", return_value="test_fp"
+    ):
+      result = await self.manager.handle_context_caching(llm_request)
 
-    # Should skip cache creation for initial request
-    assert result is None
+    # 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()

tests/unittests/flows/llm_flows/test_context_cache_processor.py

@@ -66,7 +66,7 @@ class TestContextCacheRequestProcessor:
     )
 
   def create_cache_metadata(
-      self, invocations_used=1, cache_name="test-cache", cached_contents_count=3
+      self, invocations_used=1, cache_name="test-cache", contents_count=3
   ):
     """Helper to create CacheMetadata."""
     return CacheMetadata(
@@ -76,7 +76,7 @@ class TestContextCacheRequestProcessor:
         expire_time=time.time() + 1800,
         fingerprint="test_fingerprint",
         invocations_used=invocations_used,
-        cached_contents_count=cached_contents_count,
+        contents_count=contents_count,
         created_at=time.time() - 600,
     )
 

tests/unittests/models/test_cache_metadata.py

@@ -32,7 +32,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=5,
-        cached_contents_count=3,
+        contents_count=3,
     )
 
     assert (
@@ -42,7 +42,7 @@ class TestCacheMetadata:
     assert metadata.expire_time > time.time()
     assert metadata.fingerprint == "abc123"
     assert metadata.invocations_used == 5
-    assert metadata.cached_contents_count == 3
+    assert metadata.contents_count == 3
     assert metadata.created_at is None  # Optional field
 
   def test_optional_created_at(self):
@@ -54,7 +54,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=3,
-        cached_contents_count=2,
+        contents_count=2,
         created_at=current_time,
     )
 
@@ -68,7 +68,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=0,
-        cached_contents_count=1,
+        contents_count=1,
     )
     assert metadata.invocations_used == 0
 
@@ -77,7 +77,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=10,
-        cached_contents_count=1,
+        contents_count=1,
     )
     assert metadata.invocations_used == 10
 
@@ -88,30 +88,30 @@ class TestCacheMetadata:
           expire_time=time.time() + 1800,
           fingerprint="abc123",
           invocations_used=-1,
-          cached_contents_count=1,
+          contents_count=1,
       )
     assert "greater than or equal to 0" in str(exc_info.value)
 
-  def test_cached_contents_count_validation(self):
-    """Test cached_contents_count validation constraints."""
+  def test_contents_count_validation(self):
+    """Test contents_count validation constraints."""
     # Valid: zero or positive
     metadata = CacheMetadata(
         cache_name="projects/123/locations/us-central1/cachedContents/456",
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=1,
-        cached_contents_count=0,
+        contents_count=0,
     )
-    assert metadata.cached_contents_count == 0
+    assert metadata.contents_count == 0
 
     metadata = CacheMetadata(
         cache_name="projects/123/locations/us-central1/cachedContents/456",
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=1,
-        cached_contents_count=10,
+        contents_count=10,
     )
-    assert metadata.cached_contents_count == 10
+    assert metadata.contents_count == 10
 
     # Invalid: negative
     with pytest.raises(ValidationError) as exc_info:
@@ -120,7 +120,7 @@ class TestCacheMetadata:
           expire_time=time.time() + 1800,
           fingerprint="abc123",
           invocations_used=1,
-          cached_contents_count=-1,
+          contents_count=-1,
       )
     assert "greater than or equal to 0" in str(exc_info.value)
 
@@ -133,7 +133,7 @@ class TestCacheMetadata:
         expire_time=future_time,
         fingerprint="abc123",
         invocations_used=1,
-        cached_contents_count=1,
+        contents_count=1,
     )
     assert not metadata.expire_soon
 
@@ -144,7 +144,7 @@ class TestCacheMetadata:
         expire_time=soon_time,
         fingerprint="abc123",
         invocations_used=1,
-        cached_contents_count=1,
+        contents_count=1,
     )
     assert metadata.expire_soon
 
@@ -158,7 +158,7 @@ class TestCacheMetadata:
         expire_time=expire_time,
         fingerprint="abc123",
         invocations_used=7,
-        cached_contents_count=4,
+        contents_count=4,
     )
 
     str_repr = str(metadata)
@@ -174,7 +174,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=5,
-        cached_contents_count=3,
+        contents_count=3,
     )
 
     # Should not be able to modify fields
@@ -188,7 +188,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=5,
-        cached_contents_count=3,
+        contents_count=3,
     )
 
     assert metadata.model_config["extra"] == "forbid"
@@ -201,7 +201,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=5,
-        cached_contents_count=3,
+        contents_count=3,
     )
     schema = metadata.model_json_schema()
 
@@ -211,10 +211,10 @@ class TestCacheMetadata:
         in schema["properties"]["invocations_used"]["description"]
     )
 
-    assert "cached_contents_count" in schema["properties"]
+    assert "contents_count" in schema["properties"]
     assert (
         "Number of contents"
-        in schema["properties"]["cached_contents_count"]["description"]
+        in schema["properties"]["contents_count"]["description"]
     )
 
   def test_realistic_cache_scenarios(self):
@@ -227,7 +227,7 @@ class TestCacheMetadata:
         expire_time=current_time + 1800,
         fingerprint="fresh_fingerprint",
         invocations_used=1,
-        cached_contents_count=5,
+        contents_count=5,
         created_at=current_time,
     )
     assert fresh_cache.invocations_used == 1
@@ -239,7 +239,7 @@ class TestCacheMetadata:
         expire_time=current_time + 600,
         fingerprint="used_fingerprint",
         invocations_used=8,
-        cached_contents_count=3,
+        contents_count=3,
         created_at=current_time - 1200,
     )
     assert used_cache.invocations_used == 8
@@ -252,7 +252,7 @@ class TestCacheMetadata:
         expire_time=current_time + 60,  # 1 minute
         fingerprint="expiring_fingerprint",
         invocations_used=15,
-        cached_contents_count=10,
+        contents_count=10,
     )
     assert expiring_cache.expire_soon
 
@@ -265,7 +265,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=1,
-        cached_contents_count=2,
+        contents_count=2,
     )
 
     str_repr = str(metadata)
@@ -278,7 +278,7 @@ class TestCacheMetadata:
         expire_time=time.time() + 1800,
         fingerprint="abc123",
         invocations_used=5,
-        cached_contents_count=3,
+        contents_count=3,
     )
 
     # Verify that token counts are NOT in CacheMetadata
@@ -288,22 +288,17 @@ class TestCacheMetadata:
     assert not hasattr(metadata, "prompt_tokens")
 
   def test_missing_required_fields(self):
-    """Test validation when required fields are missing."""
-    # Test each required field
+    """Test validation when truly required fields are missing."""
+    # Only fingerprint and contents_count are required now
+    # Other fields are optional (for fingerprint-only state)
     required_fields = [
-        "cache_name",
-        "expire_time",
         "fingerprint",
-        "invocations_used",
-        "cached_contents_count",
+        "contents_count",
     ]
 
     base_args = {
-        "cache_name": "projects/123/locations/us-central1/cachedContents/456",
-        "expire_time": time.time() + 1800,
         "fingerprint": "abc123",
-        "invocations_used": 1,
-        "cached_contents_count": 2,
+        "contents_count": 2,
     }
 
     for field in required_fields:
@@ -312,3 +307,13 @@ class TestCacheMetadata:
 
       with pytest.raises(ValidationError):
         CacheMetadata(**args)
+
+    # Test that optional fields can be omitted (fingerprint-only state)
+    metadata = CacheMetadata(
+        fingerprint="abc123",
+        contents_count=5,
+    )
+    assert metadata.cache_name is None
+    assert metadata.expire_time is None
+    assert metadata.invocations_used is None
+    assert metadata.created_at is None

tests/unittests/models/test_google_llm.py

@@ -96,7 +96,7 @@ def cache_metadata():
       expire_time=time.time() + 3600,
       fingerprint="test_fingerprint",
       invocations_used=2,
-      cached_contents_count=3,
+      contents_count=3,
       created_at=time.time() - 600,
   )
 

tests/unittests/utils/test_cache_performance_analyzer.py

@@ -36,7 +36,7 @@ class TestCachePerformanceAnalyzer:
     self.analyzer = CachePerformanceAnalyzer(self.mock_session_service)
 
   def create_cache_metadata(
-      self, invocations_used=1, cache_name="test-cache", cached_contents_count=5
+      self, invocations_used=1, cache_name="test-cache", contents_count=5
   ):
     """Helper to create test CacheMetadata."""
     return CacheMetadata(
@@ -46,7 +46,7 @@ class TestCachePerformanceAnalyzer:
         expire_time=time.time() + 1800,
         fingerprint="test_fingerprint",
         invocations_used=invocations_used,
-        cached_contents_count=cached_contents_count,
+        contents_count=contents_count,
         created_at=time.time() - 600,
     )