# 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