fix(vertex_ai): scope known Gemini caches by the API key actually sent

The scope used the deployment's configured Gemini key, but extra_headers can
override x-goog-api-key for the HTTP call. A caller priming a prefix with
their own key left a cache from their project that later requests with the
configured key reused and failed on. The scope now reads x-goog-api-key from
the headers sent, on the lookup, create, and single-flight paths.

litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py

@@ -3,7 +3,7 @@ import re
 import threading
 import time
 import weakref
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timezone
 from typing import Final, Literal
 
@@ -116,15 +116,17 @@ class _KnownCachedContents:
     def scope(
         base_url: str,
         custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
-        api_key: str | None,
+        headers: Mapping[str, str],
     ) -> tuple[str, str]:
         """Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project.
 
-        The Gemini key is kept as-is in this in-memory scope and is never logged or persisted.
+        The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the
+        cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted.
         """
-        if custom_llm_provider == "gemini":
-            return (base_url, api_key or "")
-        return (base_url, "")
+        if custom_llm_provider != "gemini":
+            return (base_url, "")
+        sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"]
+        return (base_url, sent_keys[-1] if sent_keys else "")
 
     @staticmethod
     def is_usable(expire_time: object) -> bool:
@@ -259,7 +261,7 @@ class ContextCachingEndpoints(VertexBase):
         self,
         cache_key: str,
         client: HTTPHandler,
-        headers: dict,
+        headers: dict[str, str],  # mutable-ok: passed through to the HTTP client, which takes a dict
         api_key: str,
         api_base: str | None,
         logging_obj: Logging,
@@ -289,7 +291,7 @@ class ContextCachingEndpoints(VertexBase):
             vertex_auth_header=vertex_auth_header,
             model=model,
         )
-        scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, api_key)
+        scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers)
         known_name: Final = known_cached_contents.get(scope, cache_key)
         if known_name is not None:
             return known_name
@@ -356,7 +358,7 @@ class ContextCachingEndpoints(VertexBase):
         self,
         cache_key: str,
         client: AsyncHTTPHandler,
-        headers: dict,
+        headers: dict[str, str],  # mutable-ok: passed through to the HTTP client, which takes a dict
         api_key: str,
         api_base: str | None,
         logging_obj: Logging,
@@ -386,7 +388,7 @@ class ContextCachingEndpoints(VertexBase):
             vertex_auth_header=vertex_auth_header,
             model=model,
         )
-        scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, api_key)
+        scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers)
         known_name: Final = known_cached_contents.get(scope, cache_key)
         if known_name is not None:
             return known_name
@@ -602,7 +604,7 @@ class ContextCachingEndpoints(VertexBase):
         raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
         _record_cache_creation_tokens(logging_obj, raw_response_cached)
         known_cached_contents.remember(
-            _KnownCachedContents.scope(url, custom_llm_provider, api_key),
+            _KnownCachedContents.scope(url, custom_llm_provider, headers),
             generated_cache_key,
             raw_response_cached.get("name"),
             raw_response_cached.get("expireTime"),
@@ -704,7 +706,7 @@ class ContextCachingEndpoints(VertexBase):
         generated_cache_key: Final = local_cache_obj.get_cache_key(
             messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
         )
-        scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, api_key)
+        scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, headers)
 
         async def lookup_or_create() -> str | None:
             google_cache_name: Final = await self.async_check_cache(

tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py

@@ -19,15 +19,18 @@ _MODULE: Final = "litellm.llms.vertex_ai.context_caching.vertex_ai_context_cachi
 _SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "")
 
 
+def _token_and_url(gemini_api_key=None, custom_llm_provider="vertex_ai", **_):
+    url: Final = "https://vertex.test/v1/projects/p/locations/us/cachedContents"
+    if custom_llm_provider == "gemini":
+        return {"x-goog-api-key": gemini_api_key}, url
+    return "token", url
+
+
 @pytest.fixture(autouse=True)
 def _cacheable_prompt_and_fixed_url():
     with (
         patch(f"{_MODULE}.is_prompt_caching_valid_prompt", return_value=True),
-        patch.object(
-            ContextCachingEndpoints,
-            "_get_token_and_url_context_caching",
-            return_value=("token", "https://vertex.test/v1/projects/p/locations/us/cachedContents"),
-        ),
+        patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", side_effect=_token_and_url),
     ):
         yield
 
@@ -53,8 +56,15 @@ def _created(name: str, expire_time: str | None) -> dict:
     return body if expire_time is None else {**body, "expireTime": expire_time}
 
 
-def _sync_call(client: HTTPHandler, api_key: str = "k", custom_llm_provider: str = "vertex_ai", messages=None):
+def _sync_call(
+    client: HTTPHandler,
+    api_key: str = "k",
+    custom_llm_provider: str = "vertex_ai",
+    messages=None,
+    extra_headers: dict | None = None,
+):
     return ContextCachingEndpoints().check_and_create_cache(
+        extra_headers=extra_headers,
         messages=messages or _messages(),
         optional_params={},
         api_key=api_key,
@@ -168,6 +178,22 @@ def test_gemini_caches_are_not_shared_across_api_keys():
     assert client.post.call_count == 2
 
 
+def test_gemini_cache_created_with_an_overridden_api_key_is_not_reused_with_the_configured_key():
+    client: Final = MagicMock(spec=HTTPHandler)
+    client.get.return_value = _response({})
+    client.post.side_effect = [
+        _response(_created("cachedContents/caller-project", _expire_time(60))),
+        _response(_created("cachedContents/deployment-project", _expire_time(60))),
+    ]
+
+    caller: Final = _sync_call(
+        client, api_key="deployment-key", custom_llm_provider="gemini", extra_headers={"X-Goog-Api-Key": "caller-key"}
+    )
+    configured: Final = _sync_call(client, api_key="deployment-key", custom_llm_provider="gemini")
+
+    assert (caller[2], configured[2]) == ("cachedContents/caller-project", "cachedContents/deployment-project")
+
+
 def test_different_prefixes_do_not_share_a_cache():
     client: Final = MagicMock(spec=HTTPHandler)
     client.get.return_value = _response({})