feat: add create_http_options to ContextCacheConfig for cache creation timeout

Merge #4702
Close #4703

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 932753541

Author: 2 people

src/google/adk/agents/context_cache_config.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from google.genai import types
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
@@ -72,6 +73,18 @@ class ContextCacheConfig(BaseModel):
       ),
   )
 
+  create_http_options: types.HttpOptions | None = Field(
+      default=None,
+      description=(
+          "Optional HTTP options to pass to the GenAI client. Set this to add a"
+          " timeout on CachedContent.create() calls (e.g."
+          " types.HttpOptions(timeout=10000) for a 10-second timeout in"
+          " milliseconds). When the cache creation call exceeds the timeout,"
+          " it fails and the request proceeds without caching. None uses the"
+          " client's default HTTP options."
+      ),
+  )
+
   @property
   def ttl_string(self) -> str:
     """Get TTL as string format for cache creation."""
@@ -81,5 +94,6 @@ def __str__(self) -> str:
     """String representation for logging."""
     return (
         f"ContextCacheConfig(cache_intervals={self.cache_intervals}, "
-        f"ttl={self.ttl_seconds}s, min_tokens={self.min_tokens})"
+        f"ttl={self.ttl_seconds}s, min_tokens={self.min_tokens}, "
+        f"create_http_options={self.create_http_options})"
     )

src/google/adk/models/gemini_context_cache_manager.py

@@ -411,6 +411,13 @@ async def _create_gemini_cache(
       if llm_request.config and llm_request.config.tool_config:
         cache_config.tool_config = llm_request.config.tool_config
 
+      # Pass through HTTP options (e.g. timeout) from cache config
+      if (
+          llm_request.cache_config
+          and llm_request.cache_config.create_http_options
+      ):
+        cache_config.http_options = llm_request.cache_config.create_http_options
+
       span.set_attribute("cache_contents_count", cache_contents_count)
       span.set_attribute("model", llm_request.model)
       span.set_attribute("ttl_seconds", llm_request.cache_config.ttl_seconds)

tests/unittests/agents/test_context_cache_config.py

@@ -106,15 +106,19 @@ def test_str_representation(self):
     )
 
     expected = (
-        "ContextCacheConfig(cache_intervals=15, ttl=3600s, min_tokens=1024)"
+        "ContextCacheConfig(cache_intervals=15, ttl=3600s, min_tokens=1024, "
+        "create_http_options=None)"
     )
     assert str(config) == expected
 
   def test_str_representation_defaults(self):
     """Test string representation with default values."""
     config = ContextCacheConfig()
 
-    expected = "ContextCacheConfig(cache_intervals=10, ttl=1800s, min_tokens=0)"
+    expected = (
+        "ContextCacheConfig(cache_intervals=10, ttl=1800s, min_tokens=0, "
+        "create_http_options=None)"
+    )
     assert str(config) == expected
 
   def test_pydantic_model_validation(self):
@@ -126,25 +130,19 @@ def test_pydantic_model_validation(self):
 
   def test_field_descriptions(self):
     """Test that fields have proper descriptions."""
-    config = ContextCacheConfig()
-    schema = config.model_json_schema()
+    fields = ContextCacheConfig.model_fields
 
-    assert "cache_intervals" in schema["properties"]
+    assert "cache_intervals" in fields
     assert (
-        "Maximum number of invocations"
-        in schema["properties"]["cache_intervals"]["description"]
+        "Maximum number of invocations" in fields["cache_intervals"].description
     )
 
-    assert "ttl_seconds" in schema["properties"]
-    assert (
-        "Time-to-live for cache"
-        in schema["properties"]["ttl_seconds"]["description"]
-    )
+    assert "ttl_seconds" in fields
+    assert "Time-to-live for cache" in fields["ttl_seconds"].description
 
-    assert "min_tokens" in schema["properties"]
+    assert "min_tokens" in fields
     assert (
-        "Minimum estimated request tokens"
-        in schema["properties"]["min_tokens"]["description"]
+        "Minimum estimated request tokens" in fields["min_tokens"].description
     )
 
   def test_immutability_config(self):

tests/unittests/agents/test_gemini_context_cache_manager.py

@@ -951,3 +951,63 @@ async def test_fingerprint_only_metadata_transitions_to_active_cache(
     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()
+
+  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

tests/unittests/test_runners.py

@@ -1304,7 +1304,8 @@ def test_runner_realistic_cache_config_scenario(self):
 
     # Verify string representation
     expected_str = (
-        "ContextCacheConfig(cache_intervals=30, ttl=14400s, min_tokens=4096)"
+        "ContextCacheConfig(cache_intervals=30, ttl=14400s, min_tokens=4096, "
+        "create_http_options=None)"
     )
     assert str(runner.context_cache_config) == expected_str