From 57fe24e2e5b60469e0e5724438f5a8db54886254 Mon Sep 17 00:00:00 2001 From: Ondra Zahradnik Date: Fri, 16 Oct 2026 10:05:52 +0200 Subject: [PATCH] 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. --- .../vertex_ai_context_caching.py | 26 +++++++------ .../test_known_cached_contents.py | 38 ++++++++++++++++--- 2 files changed, 46 insertions(+), 18 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 15df8033049..c423b0ccaf3 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/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( diff --git a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py index eb0fbb6d286..33052341756 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py +++ b/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({})