fix(vertex_ai): build the cachedContents collection URL under a custom api_base (#44783)

* fix(gemini): use collection endpoint for custom context cache

* fix(vertex_ai): build the cachedContents collection URL under a custom api_base

* fix(vertex_ai): keep a gateway prefix that contains /publishers/

* test(integration): expect the cachedContents collection URL on a custom Gemini api_base

* fix(vertex_ai): cut cache collection URLs at the base's location, graft a mount prefix

* fix(vertex_ai): cache under a publisher model api_base in that model's location

A Vertex AI api_base that names an endpoint or no model keeps its old cache URL, so
no cache is created that the chat call cannot use. The cache body takes its project and
location from the api_base model, matching the collection URL.

* fix(vertex_ai): cache for the api_base's model and skip caching on a Gemini v1 base

A Vertex api_base naming a publisher model now gives the cache its model as well as its project and location, so the cache matches the model generateContent calls. A Gemini api_base ending in /v1 skips context caching, since the v1 API has no cachedContents and rejects the cachedContent field.

* fix(vertex_ai): check the cache token minimum against the api_base's model

The cache is created for the model a publisher-model api_base names, so the minimum prompt size now comes from that model too. A prompt below its minimum is sent uncached instead of failing at cache create.

* fix(vertex_ai): use the stricter cache minimum and send a Gemini base's 404 uncached

---------

Co-authored-by: paulo037 <paulo.h.carneiro.18@gmail.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+from dataclasses import dataclass
 from typing import Final, Literal
 
 import httpx
@@ -8,7 +9,6 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
 import litellm
 from litellm._logging import verbose_logger
 from litellm.caching.caching import Cache, LiteLLMCacheType
-from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
@@ -16,10 +16,11 @@ from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
 )
 from litellm.llms.openai.openai import AllMessageValues
-from litellm.utils import is_prompt_caching_valid_prompt
+from litellm.types.llms.openai import ChatCompletionToolParam
+from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt
 
 from ..common_utils import VertexAIError, get_vertex_base_url
-from ..vertex_llm_base import VertexBase
+from ..vertex_llm_base import VertexBase, publisher_model_location
 from .known_cached_contents import KnownCachedContents, known_cached_contents
 from .transformation import (
     cached_messages_end_on_supported_turn,
@@ -65,6 +66,19 @@ class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]):
     model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
 
 
+_TOOL_PARAMS: Final = ("tools", "tool_choice")
+
+
+def _sync_cache_client(timeout: float | httpx.Timeout | None) -> HTTPHandler:
+    if timeout is None:
+        return HTTPHandler()
+    return HTTPHandler(timeout=httpx.Timeout(timeout) if isinstance(timeout, (float, int)) else timeout)
+
+
+class _CachedContentsNotServedError(VertexAIError):
+    pass
+
+
 def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedContent) -> None:
     tokens: Final = (created.get("usageMetadata") or {}).get("totalTokenCount") or 0
     logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0)  # rebind-ok: set provider count
@@ -78,6 +92,61 @@ def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None:
             return timeout
 
 
+@dataclass(frozen=True, slots=True)
+class _CacheTarget:
+    model: str
+    project: str | None
+    location: str | None
+
+
+def _cache_target(
+    custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
+    api_base: str | None,
+    model: str,
+    vertex_project: str | None,
+    vertex_location: str | None,
+) -> _CacheTarget:
+    model_location: Final = (
+        None if custom_llm_provider == "gemini" or api_base is None else publisher_model_location(api_base)
+    )
+    if model_location is None:
+        return _CacheTarget(model=model, project=vertex_project, location=vertex_location)
+    return _CacheTarget(model=model_location.model, project=model_location.project, location=model_location.location)
+
+
+def _prompt_cache_min_tokens(models: Sequence[str], custom_llm_provider: str) -> int:
+    return max(
+        get_prompt_cache_min_tokens(
+            model if model.startswith(custom_llm_provider) else f"{custom_llm_provider}/{model}"
+        )
+        for model in models
+    )
+
+
+def _context_caching_skip_reason(
+    cached_messages: Sequence[AllMessageValues],
+    tools: Sequence[ChatCompletionToolParam] | None,
+    deployment_model: str,
+    cache_model: str,
+    custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
+) -> str | None:
+    if not cached_messages_end_on_supported_turn(cached_messages):
+        return (
+            "Vertex AI context caching: cached message block ends on a model turn once "
+            "system messages are extracted, which the cachedContents API rejects."
+        )
+    min_tokens: Final = _prompt_cache_min_tokens((deployment_model, cache_model), custom_llm_provider)
+    if not is_prompt_caching_valid_prompt(
+        model=cache_model,
+        messages=list(cached_messages),
+        tools=None if tools is None else list(tools),
+        custom_llm_provider=custom_llm_provider,
+        min_token_count=min_tokens,
+    ):
+        return f"Vertex AI context caching: cached content is below minimum token count ({min_tokens})."
+    return None
+
+
 class ContextCachingEndpoints(VertexBase):
     """
     Covers context caching endpoints for Vertex AI + Google AI Studio
@@ -135,6 +204,7 @@ class ContextCachingEndpoints(VertexBase):
             vertex_project=vertex_project,
             vertex_location=vertex_location,
             vertex_api_version=("v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1"),
+            collection_endpoint=True,
         )
 
     def check_cache(
@@ -206,6 +276,8 @@ class ContextCachingEndpoints(VertexBase):
             except httpx.HTTPStatusError as e:
                 if e.response.status_code == 403:
                     return None
+                if e.response.status_code == 404 and custom_llm_provider == "gemini":
+                    raise _CachedContentsNotServedError(status_code=404, message=e.response.text)
                 raise VertexAIError(status_code=e.response.status_code, message=e.response.text)
             except Exception as e:
                 raise VertexAIError(status_code=500, message=str(e))
@@ -305,6 +377,8 @@ class ContextCachingEndpoints(VertexBase):
             except httpx.HTTPStatusError as e:
                 if e.response.status_code == 403:
                     return None
+                if e.response.status_code == 404 and custom_llm_provider == "gemini":
+                    raise _CachedContentsNotServedError(status_code=404, message=e.response.text)
                 raise VertexAIError(status_code=e.response.status_code, message=e.response.text)
             except Exception as e:
                 raise VertexAIError(status_code=500, message=str(e))
@@ -370,41 +444,25 @@ class ContextCachingEndpoints(VertexBase):
         if len(cached_messages) == 0:
             return messages, optional_params, None
 
-        if not cached_messages_end_on_supported_turn(cached_messages):
-            verbose_logger.debug(
-                "Vertex AI context caching: cached message block ends on a model turn once "
-                "system messages are extracted, which the cachedContents API rejects. "
-                "Skipping context caching."
-            )
+        target: Final = _cache_target(custom_llm_provider, api_base, model, vertex_project, vertex_location)
+        tools: Final = optional_params.get("tools")
+        tool_choice: Final = optional_params.get("tool_choice")
+        skip_reason: Final = _context_caching_skip_reason(
+            cached_messages, tools, model, target.model, custom_llm_provider
+        )
+        if skip_reason is not None:
+            verbose_logger.debug("%s Skipping context caching.", skip_reason)
             return messages, optional_params, None
 
-        # Gemini requires a minimum of 1024 tokens for context caching.
-        # Skip caching if the cached content is too small to avoid API errors.
-        if not is_prompt_caching_valid_prompt(
-            model=model,
-            messages=cached_messages,
-            tools=optional_params.get("tools"),
-            custom_llm_provider=custom_llm_provider,
-        ):
-            verbose_logger.debug(
-                "Vertex AI context caching: cached content is below minimum token "
-                "count (%d). Skipping context caching.",
-                MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
-            )
-            return messages, optional_params, None
-
-        tools: Final = optional_params.pop("tools", None)
-        tool_choice: Final = optional_params.pop("tool_choice", None)
-
         ## AUTHORIZATION ##
         token, url = self._get_token_and_url_context_caching(
             gemini_api_key=api_key,
             custom_llm_provider=custom_llm_provider,
             api_base=api_base,
-            vertex_project=vertex_project,
-            vertex_location=vertex_location,
+            vertex_project=target.project,
+            vertex_location=target.location,
             vertex_auth_header=vertex_auth_header,
-            model=model,
+            model=target.model,
         )
 
         headers: Final = {
@@ -417,44 +475,42 @@ class ContextCachingEndpoints(VertexBase):
         if extra_headers is not None:
             headers.update(extra_headers)
 
-        if client is None or not isinstance(client, HTTPHandler):
-            _params: Final = {}
-            if timeout is not None:
-                if isinstance(timeout, float) or isinstance(timeout, int):
-                    timeout = httpx.Timeout(timeout)
-                _params["timeout"] = timeout
-            client = HTTPHandler(**_params)
-        else:
-            client = client
+        http_client: Final = client if isinstance(client, HTTPHandler) else _sync_cache_client(timeout)
 
         ## CHECK IF CACHED ALREADY
         generated_cache_key: Final = local_cache_obj.get_cache_key(
-            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
-        )
-        google_cache_name: Final = self.check_cache(
-            cache_key=generated_cache_key,
-            client=client,
-            headers=headers,
-            api_key=api_key,
-            api_base=api_base,
-            logging_obj=logging_obj,
-            custom_llm_provider=custom_llm_provider,
-            vertex_project=vertex_project,
-            vertex_location=vertex_location,
-            vertex_auth_header=vertex_auth_header,
-            model=model,
+            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=target.model
         )
+        try:
+            google_cache_name: Final = self.check_cache(
+                cache_key=generated_cache_key,
+                client=http_client,
+                headers=headers,
+                api_key=api_key,
+                api_base=api_base,
+                logging_obj=logging_obj,
+                custom_llm_provider=custom_llm_provider,
+                vertex_project=target.project,
+                vertex_location=target.location,
+                vertex_auth_header=vertex_auth_header,
+                model=target.model,
+            )
+        except _CachedContentsNotServedError:
+            verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url)
+            return messages, optional_params, None
+        for tool_param in _TOOL_PARAMS:
+            optional_params.pop(tool_param, None)
         if google_cache_name:
             return non_cached_messages, optional_params, google_cache_name
 
         ## TRANSFORM REQUEST
         cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching(
-            model=model,
+            model=target.model,
             messages=cached_messages,
             cache_key=generated_cache_key,
             custom_llm_provider=custom_llm_provider,
-            vertex_project=vertex_project,
-            vertex_location=vertex_location,
+            vertex_project=target.project,
+            vertex_location=target.location,
         )
 
         cached_content_request_body["tools"] = tools
@@ -473,7 +529,7 @@ class ContextCachingEndpoints(VertexBase):
         )
 
         try:
-            response: Final = client.post(
+            response: Final = http_client.post(
                 url=url,
                 headers=headers,
                 json=cached_content_request_body,
@@ -489,7 +545,7 @@ class ContextCachingEndpoints(VertexBase):
         _record_cache_creation_tokens(logging_obj, raw_response_cached)
         self._known_caches.remember(
             KnownCachedContents.scope(
-                url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
+                url, custom_llm_provider, headers, vertex_project=target.project, vertex_location=target.location
             ),
             generated_cache_key,
             raw_response_cached.get("name"),
@@ -536,41 +592,25 @@ class ContextCachingEndpoints(VertexBase):
         if len(cached_messages) == 0:
             return messages, optional_params, None
 
-        if not cached_messages_end_on_supported_turn(cached_messages):
-            verbose_logger.debug(
-                "Vertex AI context caching: cached message block ends on a model turn once "
-                "system messages are extracted, which the cachedContents API rejects. "
-                "Skipping context caching."
-            )
+        target: Final = _cache_target(custom_llm_provider, api_base, model, vertex_project, vertex_location)
+        tools: Final = optional_params.get("tools")
+        tool_choice: Final = optional_params.get("tool_choice")
+        skip_reason: Final = _context_caching_skip_reason(
+            cached_messages, tools, model, target.model, custom_llm_provider
+        )
+        if skip_reason is not None:
+            verbose_logger.debug("%s Skipping context caching.", skip_reason)
             return messages, optional_params, None
 
-        # Gemini requires a minimum of 1024 tokens for context caching.
-        # Skip caching if the cached content is too small to avoid API errors.
-        if not is_prompt_caching_valid_prompt(
-            model=model,
-            messages=cached_messages,
-            tools=optional_params.get("tools"),
-            custom_llm_provider=custom_llm_provider,
-        ):
-            verbose_logger.debug(
-                "Vertex AI context caching: cached content is below minimum token "
-                "count (%d). Skipping context caching.",
-                MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
-            )
-            return messages, optional_params, None
-
-        tools: Final = optional_params.pop("tools", None)
-        tool_choice: Final = optional_params.pop("tool_choice", None)
-
         ## AUTHORIZATION ##
         token, url = self._get_token_and_url_context_caching(
             gemini_api_key=api_key,
             custom_llm_provider=custom_llm_provider,
             api_base=api_base,
-            vertex_project=vertex_project,
-            vertex_location=vertex_location,
+            vertex_project=target.project,
+            vertex_location=target.location,
             vertex_auth_header=vertex_auth_header,
-            model=model,
+            model=target.model,
         )
 
         headers: Final = {
@@ -590,10 +630,10 @@ class ContextCachingEndpoints(VertexBase):
 
         ## CHECK IF CACHED ALREADY
         generated_cache_key: Final = local_cache_obj.get_cache_key(
-            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
+            messages=cached_messages, tools=tools, tool_choice=tool_choice, model=target.model
         )
         scope: Final = KnownCachedContents.scope(
-            url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
+            url, custom_llm_provider, headers, vertex_project=target.project, vertex_location=target.location
         )
 
         async def lookup_or_create() -> str | None:
@@ -605,10 +645,10 @@ class ContextCachingEndpoints(VertexBase):
                 api_base=api_base,
                 logging_obj=logging_obj,
                 custom_llm_provider=custom_llm_provider,
-                vertex_project=vertex_project,
-                vertex_location=vertex_location,
+                vertex_project=target.project,
+                vertex_location=target.location,
                 vertex_auth_header=vertex_auth_header,
-                model=model,
+                model=target.model,
             )
 
             if google_cache_name:
@@ -616,12 +656,12 @@ class ContextCachingEndpoints(VertexBase):
 
             ## TRANSFORM REQUEST
             cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching(
-                model=model,
+                model=target.model,
                 messages=cached_messages,
                 cache_key=generated_cache_key,
                 custom_llm_provider=custom_llm_provider,
-                vertex_project=vertex_project,
-                vertex_location=vertex_location,
+                vertex_project=target.project,
+                vertex_location=target.location,
             )
 
             cached_content_request_body["tools"] = tools
@@ -659,9 +699,15 @@ class ContextCachingEndpoints(VertexBase):
             )
             return raw_response_cached.get("name")
 
-        cache_name: Final = await self._known_caches.single_flight(
-            scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout)
-        )
+        try:
+            cache_name: Final = await self._known_caches.single_flight(
+                scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout)
+            )
+        except _CachedContentsNotServedError:
+            verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url)
+            return messages, optional_params, None
+        for tool_param in _TOOL_PARAMS:
+            optional_params.pop(tool_param, None)
         return non_cached_messages, optional_params, cache_name
 
     def get_cache(self):

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -6,9 +6,11 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
 
 import asyncio
 import json
+import re
 import threading
 from collections import OrderedDict
 from collections.abc import Mapping
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
 from urllib.parse import urlparse
 
@@ -46,6 +48,31 @@ def graft_default_vertex_path(api_base: str, default_url: str) -> str:
 
 _graft_default_vertex_path = graft_default_vertex_path
 
+
+_PUBLISHER_MODEL_LOCATION: Final = re.compile(r".*/projects/[^/]+/locations/[^/]+(?=/publishers/[^/]+/models/[^/]+/?$)")
+
+
+@dataclass(frozen=True, slots=True)
+class PublisherModelLocation:
+    parent_url: str
+    project: str
+    location: str
+    model: str
+
+
+def publisher_model_location(api_base: str) -> PublisherModelLocation | None:
+    match: Final = _PUBLISHER_MODEL_LOCATION.match(api_base)
+    if match is None:
+        return None
+    _, _, project, _, location = match.group().rsplit("/", 4)
+    return PublisherModelLocation(
+        parent_url=match.group(),
+        project=project,
+        location=location,
+        model=api_base.rstrip("/").rsplit("/", 1)[-1],
+    )
+
+
 GOOGLE_IMPORT_ERROR_MESSAGE: Final = (
     "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform"
 )
@@ -637,15 +664,18 @@ class VertexBase:
         vertex_location: str | None = None,
         vertex_api_version: Literal["v1", "v1beta1"] | None = None,
         use_psc_endpoint_format: bool = False,
+        collection_endpoint: bool = False,
     ) -> tuple[str | None, str]:
         """
         for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317
 
         Handles custom api_base for:
-        1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}
+        1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}, or {api_base}/{endpoint}
+           for a collection endpoint
         2. Vertex AI with standard proxies - grafts the default vertex URL path onto the
            api_base when its path is empty or only an API version (/v1, /v1beta1);
-           otherwise constructs {api_base}:{endpoint}
+           otherwise constructs {api_base}:{endpoint}, or, for a collection endpoint under an api_base
+           naming a publisher model, the collection under that model's projects/{project}/locations/{location}
         3. Vertex AI with PSC endpoints - constructs full path structure
            {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}
            (only when use_psc_endpoint_format=True)
@@ -653,22 +683,22 @@ class VertexBase:
         Args:
             use_psc_endpoint_format: If True, constructs PSC endpoint URL format.
                                      If False (default), uses api_base as-is and appends :{endpoint}
+            collection_endpoint: True when `endpoint` names a collection resource such as
+                                 cachedContents rather than a model action such as generateContent
 
         ## Returns
         - (auth_header, url) - Tuple[Optional[str], str]
         """
         if api_base:
             if custom_llm_provider == "gemini":
-                # For Gemini (Google AI Studio), construct the full path like other providers
-                if model is None:
+                if model is None and not collection_endpoint:
                     raise ValueError("Model parameter is required for Gemini custom API base URLs")
-                url = join_api_base(api_base, f"models/{model}:{endpoint}")
+                url = join_api_base(api_base, endpoint if collection_endpoint else f"models/{model}:{endpoint}")
                 if gemini_api_key is None:
                     raise ValueError(
                         "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
                     )
-                if gemini_api_key is not None:
-                    auth_header = {"x-goog-api-key": gemini_api_key}
+                auth_header = {"x-goog-api-key": gemini_api_key}
             else:
                 # For Vertex AI
                 if use_psc_endpoint_format:
@@ -692,7 +722,12 @@ class VertexBase:
                 elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path:
                     url = graft_default_vertex_path(api_base=api_base, default_url=url)
                 else:
-                    url = f"{api_base.rstrip('/')}:{endpoint}"
+                    model_location: Final = publisher_model_location(api_base) if collection_endpoint else None
+                    url = (
+                        f"{api_base.rstrip('/')}:{endpoint}"
+                        if model_location is None
+                        else join_api_base(model_location.parent_url, endpoint)
+                    )
             if stream is True:
                 parsed_stream_url: Final = urlparse(url)
                 stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse"

litellm/types/llms/vertex_ai.py

@@ -376,7 +376,7 @@ class RequestBody(TypedDict, total=False):
 class CachedContentRequestBody(TypedDict, total=False):
     contents: Required[list[ContentType]]
     system_instruction: SystemInstructions
-    tools: Tools
+    tools: Tools | None
     toolConfig: ToolConfig
     model: Required[str]  # Format: models/{model}
     ttl: str  # ending in 's' - Example: "3.5s".

tests/integration/cost_calculation/cases/gemini.py

@@ -3348,10 +3348,10 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase(
     response=RoutedResponse(
         content_type="application/x-routed",
         routes={
-            "GET /models/gemini-3.1-pro:cachedContents": JsonResponse(
+            "GET /cachedContents": JsonResponse(
                 content_type="application/json", body={"cachedContents": []}
             ),
-            "POST /models/gemini-3.1-pro:cachedContents": JsonResponse(
+            "POST /cachedContents": JsonResponse(
                 content_type="application/json",
                 body={
                     "name": "cachedContents/$REQUEST_ID",

tests/integration/messages_endpoint/providers/gemini/test_gemini_messages_cache_control_wire.py

@@ -39,11 +39,11 @@ def test_gemini_messages_cache_control_creates_cached_content_and_generates_from
     def respond(request: Request) -> Reply:
         assert request.headers["x-goog-api-key"] == _API_KEY, request.headers
         if request.method == "GET":
-            assert request.target == f"/models/{_BACKEND}:cachedContents", request.target
+            assert request.target == "/cachedContents", request.target
             return Reply(body=b"{}")
         assert request.method == "POST", request.method
         body: Final = _JSON_OBJECT.validate_json(request.body)
-        if request.target == f"/models/{_BACKEND}:cachedContents":
+        if request.target == "/cachedContents":
             assert isinstance(body["displayName"], str) and body["displayName"], body
             assert body == {
                 "contents": [{"role": "user", "parts": [{"text": "."}]}],
@@ -80,7 +80,7 @@ def test_gemini_messages_cache_control_creates_cached_content_and_generates_from
         payload: Final = _JSON_OBJECT.validate_json(response.content)
         assert payload["content"] == [{"type": "text", "text": "The policy applies."}], response.text
         assert [(request.method, request.target) for request in wire.drain()] == [
-            ("GET", f"/models/{_BACKEND}:cachedContents"),
-            ("POST", f"/models/{_BACKEND}:cachedContents"),
+            ("GET", "/cachedContents"),
+            ("POST", "/cachedContents"),
             ("POST", f"/models/{_BACKEND}:generateContent"),
         ]

tests/integration/providers/test_cache_control_tool_call_marks_wire.py

@@ -406,8 +406,8 @@ def _gemini_peer(request: Request) -> Reply:
 
 _GEMINI: Final = "gemini/gemini-3.8-flash"
 _GEMINI_CACHE_WRITE: Final = [
-    ("GET", "/models/gemini-3.8-flash:cachedContents"),
-    ("POST", "/models/gemini-3.8-flash:cachedContents"),
+    ("GET", "/cachedContents"),
+    ("POST", "/cachedContents"),
     ("POST", "/models/gemini-3.8-flash:generateContent"),
 ]
 

tests/integration/providers/test_gemini_context_caching_custom_api_base_wire.py

@@ -0,0 +1,130 @@
+import json
+from collections.abc import Callable
+from typing import Final
+
+import pytest
+
+from integration._support.client import Gateway
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_BACKEND: Final = "gemini-3.8-flash"
+_PROJECT: Final = "scripted-project"
+_LOCATION: Final = "us-central1"
+_MODEL_LOCATION: Final = "europe-west4"
+_VERTEX_MODEL_RESOURCE: Final = f"projects/{_PROJECT}/locations/{_MODEL_LOCATION}/publishers/google/models/{_BACKEND}"
+_VERTEX_COLLECTION: Final = f"/v1/projects/{_PROJECT}/locations/{_MODEL_LOCATION}/cachedContents"
+_GEMINI_COLLECTION: Final = "/v1beta/cachedContents"
+_GEMINI_MODEL_PATH: Final = f"/v1beta/models/{_BACKEND}"
+_GEMINI_KEY: Final = "scripted-gemini-key"
+_CACHE_NAME: Final = "cachedContents/scripted-cache"
+_CACHED_TOKENS: Final = 5120
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_CACHED_TEXT: Final = " ".join(
+    f"Policy clause {index}: the gateway keeps this clause in the cache." for index in range(700)
+)
+_MESSAGES: Final = (
+    {
+        "role": "user",
+        "content": [{"type": "text", "text": _CACHED_TEXT, "cache_control": {"type": "ephemeral"}}],
+    },
+    {"role": "user", "content": "How many clauses are there?"},
+)
+
+
+def _generate_content_reply() -> Reply:
+    return Reply(
+        body=json.dumps(
+            {
+                "candidates": [{"content": {"role": "model", "parts": [{"text": "700"}]}, "finishReason": "STOP"}],
+                "usageMetadata": {
+                    "promptTokenCount": _CACHED_TOKENS + 9,
+                    "cachedContentTokenCount": _CACHED_TOKENS,
+                    "candidatesTokenCount": 2,
+                    "totalTokenCount": _CACHED_TOKENS + 11,
+                },
+                "modelVersion": _BACKEND,
+            }
+        ).encode()
+    )
+
+
+def _google_responder(collection: str, model_path: str) -> Callable[[Request], Reply]:
+    def respond(request: Request) -> Reply:
+        if request.target == collection and request.method == "GET":
+            return Reply(body=b'{"cachedContents": []}')
+        if request.target == collection and request.method == "POST":
+            return Reply(body=json.dumps({"name": _CACHE_NAME, "model": f"models/{_BACKEND}"}).encode())
+        assert request.target == f"{model_path}:generateContent", f"{request.method} {request.target}"
+        return _generate_content_reply()
+
+    return respond
+
+
+def _chat(gateway: Gateway, model: str) -> dict[str, JsonValue]:
+    response: Final = gateway.client.post(
+        "/v1/chat/completions",
+        json={"model": model, "messages": list(_MESSAGES), "max_tokens": 32},
+        headers={"Authorization": f"Bearer {gateway.key}"},
+        timeout=30,
+    )
+    assert response.status_code == 200, response.text
+    return _JSON_OBJECT.validate_json(response.content)
+
+
+def _assert_cache_round_trip(requests: tuple[Request, ...], collection: str, model_path: str) -> None:
+    assert tuple((request.method, request.target) for request in requests) == (
+        ("GET", collection),
+        ("POST", collection),
+        ("POST", f"{model_path}:generateContent"),
+    ), [(request.method, request.target) for request in requests]
+    created: Final = _JSON_OBJECT.validate_json(requests[1].body)
+    assert "contents" in created, created
+    generated: Final = _JSON_OBJECT.validate_json(requests[2].body)
+    assert generated["cachedContent"] == _CACHE_NAME, generated
+
+
+def _cached_tokens_of(response: dict[str, JsonValue]) -> JsonValue:
+    usage: Final = response["usage"]
+    assert isinstance(usage, dict)
+    details: Final = usage["prompt_tokens_details"]
+    assert isinstance(details, dict)
+    return details["cached_tokens"]
+
+
+def test_gemini_custom_api_base_lists_and_creates_the_cached_contents_collection(gateway: Gateway) -> None:
+    with (
+        wire_server(_google_responder(_GEMINI_COLLECTION, _GEMINI_MODEL_PATH)) as wire,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=f"{wire.url}/v1beta", api_key=_GEMINI_KEY)
+        response: Final = _chat(gateway, model)
+        requests: Final = wire.drain()
+        _assert_cache_round_trip(requests, _GEMINI_COLLECTION, _GEMINI_MODEL_PATH)
+        assert all(request.headers.get("x-goog-api-key") == _GEMINI_KEY for request in requests)
+        assert _cached_tokens_of(response) == _CACHED_TOKENS
+
+
+@pytest.mark.parametrize("gateway_prefix", ["", "/vertex-proxy"], ids=["model-path", "prefixed-model-path"])
+def test_vertex_custom_api_base_creates_the_cache_where_its_model_lives(gateway: Gateway, gateway_prefix: str) -> None:
+    api_base_path: Final = f"{gateway_prefix}/v1/{_VERTEX_MODEL_RESOURCE}"
+    collection: Final = f"{gateway_prefix}{_VERTEX_COLLECTION}"
+    with (
+        wire_server(_google_responder(collection, api_base_path)) as wire,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(
+            model=f"vertex_ai/{_BACKEND}",
+            api_base=f"{wire.url}{api_base_path}",
+            api_key=None,
+            vertex_project=_PROJECT,
+            vertex_location=_LOCATION,
+            vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
+        )
+        response: Final = _chat(gateway, model)
+        requests: Final = wire.drain()
+        _assert_cache_round_trip(requests, collection, api_base_path)
+        assert _JSON_OBJECT.validate_json(requests[1].body)["model"] == _VERTEX_MODEL_RESOURCE
+        assert all(request.headers.get("authorization") == "Bearer scripted-token" for request in requests)
+        assert _cached_tokens_of(response) == _CACHED_TOKENS