fix(models): gate Gemini cache creation on cacheable prefix tokens

Merge google#6137

### Link to Issue or Description of Change

- Closes: google#5847

**Problem:**

`GeminiContextCacheManager._create_new_cache_with_contents` decides whether a
request is large enough to cache by comparing
`llm_request.cacheable_contents_token_count` against Gemini's
`_GEMINI_MIN_CACHE_TOKENS` (4096). But `cacheable_contents_token_count` is set
(in `context_cache_processor.py`) to the **previous response's
`prompt_token_count`** — i.e. the token count of the *entire* previous prompt
(system instruction + tools + every content, including the trailing user turn).

The cache, however, only stores the **prefix** `contents[:cache_contents_count]`
(plus system instruction and tools — see `_create_gemini_cache`), where
`cache_contents_count` excludes the last continuous batch of user contents.

On a long conversation the full-prompt count can clear the 4096 gate while the
actual cached prefix is far below it, so `caches.create()` is invoked with a
sub-minimum payload and Gemini returns **`400 INVALID_ARGUMENT`**. The reporter
confirmed this in production, and a maintainer confirmed the root cause in the
issue thread.

**Solution:**

Gate cache creation on the size of the prefix that is *actually* cached rather
than on the full previous prompt.

Since the only accurate token count we have is for the whole prompt, the new
`_estimate_cacheable_prefix_tokens` scales that accurate count by the prefix's
estimated share of the request (`_estimate_request_tokens` is reused — now
accepting an optional `cache_contents_count` — to estimate both the full request
and the prefix). When the prefix already spans the whole request the scale
factor is `1.0` and the accurate count is used unchanged, so existing behavior
is preserved; when the prefix is a strict subset the estimate shrinks
accordingly and a too-small prefix is correctly skipped before any API call.

Scope note: this fixes the `400 INVALID_ARGUMENT` (the reported bug). The
separate latency concern discussed later in the thread (synchronous
`caches.create()` / cleanup on the request path, tracked in google#5889) is
intentionally out of scope. The user-configured `min_tokens` gate is left on the
full-prompt count, since the 4096 prefix gate is what prevents the invalid API
call; happy to extend `min_tokens` to the prefix too if maintainers prefer.

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

Added `test_create_cache_gates_on_prefix_not_full_prompt`: a tiny cacheable
prefix followed by a huge trailing user turn with
`cacheable_contents_token_count = 75000`. It asserts no cache is created and
`caches.create` is never called. This test is **red on `main`** (`caches.create`
is called once — the 400-inducing path) and **green with this change**.

Updated `test_fingerprint_only_metadata_transitions_to_active_cache` to use a
realistically large cacheable prefix, since its previous setup relied on the
buggy assumption that a tiny prefix would still be cached because the injected
full-prompt count cleared the gate.

```
$ pytest tests/unittests/agents/test_gemini_context_cache_manager.py \
         tests/unittests/agents/test_context_cache_config.py \
         tests/unittests/flows/llm_flows/test_context_cache_processor.py -q
59 passed
```

**Manual End-to-End (E2E) Tests:**

The fix is exercised entirely through the unit test above (it mocks
`genai_client.aio.caches.create` and asserts it is not called for a
below-minimum prefix), which mirrors the production failure mode reported in
google#5847 (`400 INVALID_ARGUMENT` from `caches.create` on a long conversation).

### Checklist

- [x] I have read the CONTRIBUTING.md document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective.
- [x] New and existing unit tests pass locally with my changes.

Co-authored-by: George Weale <gweale@google.com>
COPYBARA_INTEGRATE_REVIEW=google#6137 from nyxst4ck:fix/gemini-cache-prefix-token-gate ade7a1e
PiperOrigin-RevId: 938659723

Author: 2 people

src/google/adk/models/gemini_context_cache_manager.py

@@ -326,11 +326,21 @@ 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:
+    # `cacheable_contents_token_count` is the token count of the whole previous
+    # prompt (system instruction + tools + every content). The cache, however,
+    # only stores the prefix `contents[:cache_contents_count]` plus the system
+    # instruction and tools (see `_create_gemini_cache`). On a long conversation
+    # the full-prompt count can clear Gemini's minimum while the cached prefix
+    # is far smaller, which makes `caches.create` fail with 400
+    # INVALID_ARGUMENT.
+    # Gate on the estimated prefix size so we never send a sub-minimum payload.
+    cacheable_prefix_tokens = self._estimate_cacheable_prefix_tokens(
+        llm_request, cache_contents_count
+    )
+    if cacheable_prefix_tokens < _GEMINI_MIN_CACHE_TOKENS:
       logger.info(
-          "Request below Gemini minimum cache size (%d < %d tokens)",
-          llm_request.cacheable_contents_token_count,
+          "Cacheable prefix below Gemini minimum cache size (%d < %d tokens)",
+          cacheable_prefix_tokens,
           _GEMINI_MIN_CACHE_TOKENS,
       )
       return None
@@ -342,13 +352,20 @@ async def _create_new_cache_with_contents(
       logger.warning("Failed to create cache: %s", e)
       return None
 
-  def _estimate_request_tokens(self, llm_request: LlmRequest) -> int:
-    """Estimate token count for the request.
+  def _estimate_request_tokens(
+      self,
+      llm_request: LlmRequest,
+      cache_contents_count: Optional[int] = None,
+  ) -> int:
+    """Estimate token count for the request (or its cacheable prefix).
 
     This is a rough estimation based on content text length.
 
     Args:
         llm_request: Request to estimate tokens for
+        cache_contents_count: When provided, only the first
+            ``cache_contents_count`` contents are counted (the prefix that gets
+            cached); the system instruction and tools are always included.
 
     Returns:
         Estimated token count
@@ -366,15 +383,54 @@ def _estimate_request_tokens(self, llm_request: LlmRequest) -> int:
           tool_str = json.dumps(tool.model_dump())
           total_chars += len(tool_str)
 
-    # Contents
-    for content in llm_request.contents:
+    # Contents (optionally limited to the cacheable prefix)
+    contents = llm_request.contents
+    if cache_contents_count is not None:
+      contents = contents[:cache_contents_count]
+    for content in contents:
       for part in content.parts:
         if part.text:
           total_chars += len(part.text)
 
     # Rough estimate: 4 characters per token
     return total_chars // 4
 
+  def _estimate_cacheable_prefix_tokens(
+      self, llm_request: LlmRequest, cache_contents_count: int
+  ) -> int:
+    """Estimate the token count of the prefix that will actually be cached.
+
+    The only accurate token count available is
+    ``cacheable_contents_token_count``, which covers the entire previous prompt.
+    Since the cache stores just the prefix ``contents[:cache_contents_count]``
+    (plus system instruction and tools), we scale that accurate count by the
+    prefix's estimated share of the request. When the prefix already spans the
+    whole request the scale factor is 1 and the accurate count is returned
+    unchanged.
+
+    Args:
+        llm_request: Request to estimate the cacheable prefix tokens for
+        cache_contents_count: Number of leading contents that get cached
+
+    Returns:
+        Estimated token count of the cacheable prefix
+    """
+    full_tokens = llm_request.cacheable_contents_token_count
+    if not full_tokens:
+      return 0
+
+    full_estimate = self._estimate_request_tokens(llm_request)
+    if full_estimate <= 0:
+      # No text to estimate from (e.g. non-text parts); fall back to the
+      # accurate full count rather than incorrectly skipping the cache.
+      return full_tokens
+
+    prefix_estimate = self._estimate_request_tokens(
+        llm_request, cache_contents_count
+    )
+    ratio = min(1.0, prefix_estimate / full_estimate)
+    return int(full_tokens * ratio)
+
   async def _create_gemini_cache(
       self, llm_request: LlmRequest, cache_contents_count: int
   ) -> CacheMetadata:

tests/unittests/agents/test_gemini_context_cache_manager.py

@@ -202,6 +202,43 @@ async def test_handle_context_caching_invalid_cache_fingerprint_match(self):
     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,
   ):
@@ -916,7 +953,10 @@ async def test_fingerprint_only_metadata_transitions_to_active_cache(
     llm_request_2 = self.create_llm_request(
         cache_metadata=result_1, contents_count=5
     )
-    llm_request_2.cacheable_contents_token_count = 4096
+    # 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),
@@ -997,7 +1037,10 @@ async def test_dynamic_instruction_does_not_break_initial_cache_fingerprint(
         dynamic_instruction,
         tool_response,
     ]
-    request_2.cacheable_contents_token_count = 4096
+    # 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 = (