fix(vertex_ai): reuse one Google token across RAG Engine searches and cold bursts (#51239)

* fix(vertex_ai): reuse one Google token across vector store searches

* fix(vertex_ai): load Vertex credentials once when a cold worker gets a burst

* test(vertex_ai): count Google token fetches for RAG Engine searches through the proxy

* test(vertex_ai): type the fixtures and token callback in the concurrent search test

* fix(vertex_ai): lock Google token loads per credential and reload a key Google rejects

A shared RAG Engine config made one store's slow token fetch hold up every
other store, and kept a rotated key file's deleted key until a restart.
First loads and expired refreshes now take a lock per credentials and project,
refreshes take a lock per credentials object, and any RefreshError reloads the
credentials once.

* fix(vertex_ai): share one Google token load or failure per credential and reload rejected keys on async calls

* test(vertex_ai): check async callers sharing an expired credential refresh it once

* fix(vertex_ai): reload Google credentials only when a refresh is rejected for good

A retryable refresh error such as a 503 now raises without reloading the key, so an outage
does not double google-auth's retries. Credentials that cannot be weak-referenced share the
config's refresh lock again, as on main. The burst test for a rejected key moves to VertexBase
and holds the token response until the other callers have started

* fix(vertex_ai): reload Google credentials only when Google answers invalid_grant

google-auth leaves its retryable flag False on metadata server, impersonation, and workload identity outages, and on a 502, so the previous check reloaded the key file on those outages and doubled the token tries. Now only an invalid_grant answer or "Reauthentication is needed" reloads

The sync reload also refills the cache under the key the caller asked with, so a call without vertex_project keeps one cached copy, as the async reload already did

* fix(vertex_ai): pick up a swapped Google key file on the next call

The credential cache key now carries the key file's inode, mtime and size, for a
vertex_credentials path and for GOOGLE_APPLICATION_CREDENTIALS. A key file that is replaced
or rewritten, such as a rotated Kubernetes secret, is loaded on the next call instead of
after Google rejects the old key or the proxy restarts

* fix(vertex_ai): keep the earlier key while a replaced key file does not load

A key file caught halfway through a non-atomic write, or replaced by one that does not load,
now keeps the credentials loaded from its earlier version in use and logs a warning, as the
cache did before key files were re-read. The next complete file is loaded on the next call.
The pub/sub test no longer writes a key file it only reaches through the seeded cache

* fix(vertex_ai): try a replaced key file again when its first token fetch fails

The earlier key now stands in only when the replaced file does not load. A token fetch that fails on the new key raises like any first load, so the next call tries the new key again instead of keeping the old one

* fix(vertex): raise when an async reload keeps an earlier key Google rejected

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/llms/vertex_ai/credentials_source.py

@@ -159,6 +159,14 @@ def load_vertex_credentials_source(credentials: str) -> VertexCredentialsSource:
     return _read_json_file(credentials)
 
 
+def vertex_credentials_file_version(path: str) -> str | None:
+    try:
+        stat: Final = os.stat(path)
+    except (OSError, ValueError):
+        return None
+    return f"{stat.st_ino}:{stat.st_mtime_ns}:{stat.st_size}"
+
+
 def build_vertex_credentials(
     source: VertexCredentialsJson, build: Callable[[Mapping[str, object]], _Credentials]
 ) -> _Credentials:

litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py

@@ -361,3 +361,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
                 status_code=response.status_code,
                 headers=response.headers,
             )
+
+
+vertex_rag_api_vector_store_config: Final = VertexVectorStoreConfig()

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -6,10 +6,13 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
 
 import asyncio
 import json
+import os
 import re
 import threading
+import weakref
 from collections import OrderedDict
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
+from concurrent.futures import Future
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
 from urllib.parse import urlparse
@@ -35,6 +38,7 @@ from .credentials_source import (
     build_vertex_credentials,
     load_vertex_credentials_source,
     raise_vertex_credentials_failure,
+    vertex_credentials_file_version,
 )
 
 
@@ -100,6 +104,32 @@ class _VertexCredentialsObject(Protocol):
 
 
 _VertexCredentialCacheKey = tuple[str | None, str | None]
+_KEY_FILE_VERSION_SEPARATOR: Final = "\0"
+
+
+def _credential_cache_identity(credentials: VERTEX_CREDENTIALS_TYPES | None) -> str | None:
+    if isinstance(credentials, dict):
+        return json.dumps(credentials)
+    key_file: Final = credentials if credentials is not None else os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
+    version: Final = None if key_file is None else vertex_credentials_file_version(key_file)
+    if version is None:
+        return credentials
+    return f"{credentials or ''}{_KEY_FILE_VERSION_SEPARATOR}{version}"
+
+
+def _is_other_version_of_key_file(
+    candidate: _VertexCredentialCacheKey, credential_cache_key: _VertexCredentialCacheKey
+) -> bool:
+    candidate_identity, candidate_project_id = candidate
+    identity, project_id = credential_cache_key
+    if candidate_identity is None or identity is None or candidate_identity == identity:
+        return False
+    key_file, separator, _ = candidate_identity.partition(_KEY_FILE_VERSION_SEPARATOR)
+    return (
+        bool(separator)
+        and candidate_project_id == project_id
+        and key_file == identity.partition(_KEY_FILE_VERSION_SEPARATOR)[0]
+    )
 
 
 class VertexBase:
@@ -128,11 +158,13 @@ class VertexBase:
         self._async_refresh_lock_refcounts: dict[_VertexCredentialCacheKey, int] = {}
         # Tracks in-flight background refresh tasks to avoid duplicate refreshes.
         self._background_refresh_tasks: dict[_VertexCredentialCacheKey, asyncio.Task] = {}
-        # Protects the sync get_access_token refresh path.
-        # Use RLock so that the reauthentication retry path (which calls
-        # back into get_access_token while still holding the lock) can
-        # re-acquire it without deadlocking the current thread.
-        self._sync_refresh_lock = threading.RLock()
+        self._sync_token_flights: weakref.WeakValueDictionary[_VertexCredentialCacheKey, Future[tuple[str, str]]] = (
+            weakref.WeakValueDictionary()
+        )
+        self._credential_refresh_locks: weakref.WeakKeyDictionary[_VertexCredentialsObject, threading.RLock] = (
+            weakref.WeakKeyDictionary()
+        )
+        self._unkeyable_credentials_refresh_lock: Final = threading.RLock()
 
     @staticmethod
     def get_vertex_region(vertex_region: str | None, model: str) -> str:
@@ -399,18 +431,64 @@ class VertexBase:
         except ImportError:
             raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE)
 
-        # Serialize all refreshes on this VertexBase across threads.
         # ``credentials.refresh()`` is not safe to call concurrently on the
         # same credentials object, and this method is invoked from three
         # places that can run on different threads:
-        #   - sync ``get_access_token`` (already holds ``_sync_refresh_lock``)
+        #   - sync ``get_access_token``
         #   - the async slow path (via ``asyncify`` in a worker thread)
         #   - the background proactive refresh task (via ``asyncify``)
-        # ``_sync_refresh_lock`` is an ``RLock`` so reentrant acquisition
-        # from the sync path is safe.
-        with self._sync_refresh_lock:
+        with self._credential_refresh_lock(credentials):
             credentials.refresh(Request())
 
+    def _credential_refresh_lock(self, credentials: _VertexCredentialsObject) -> threading.RLock:
+        with self._credentials_cache_lock:
+            try:
+                return self._credential_refresh_locks.setdefault(credentials, threading.RLock())
+            except TypeError:
+                return self._unkeyable_credentials_refresh_lock
+
+    def _refresh_if_still_needed(
+        self,
+        credentials: _VertexCredentialsObject,
+        needs_refresh: Callable[[_VertexCredentialsObject], object],
+    ) -> None:
+        with self._credential_refresh_lock(credentials):
+            if needs_refresh(credentials):
+                verbose_logger.debug("Credentials expired, refreshing")
+                self.refresh_auth(credentials)
+
+    def _needs_blocking_refresh(self, credentials: _VertexCredentialsObject) -> bool:
+        from google.auth.credentials import TokenState
+
+        return not isinstance(credentials.token, str) or self._get_token_state(credentials) == TokenState.INVALID
+
+    def _join_sync_token_flight(
+        self, credential_cache_key: _VertexCredentialCacheKey
+    ) -> tuple[Future[tuple[str, str]], bool]:
+        with self._credentials_cache_lock:
+            in_flight: Final = self._sync_token_flights.get(credential_cache_key)
+            if in_flight is not None:
+                return in_flight, False
+            flight: Final[Future[tuple[str, str]]] = Future()
+            self._sync_token_flights[credential_cache_key] = flight
+            return flight, True
+
+    @staticmethod
+    def _refresh_rejected(error: Exception) -> bool:
+        from google.auth.exceptions import RefreshError
+
+        match error:
+            case RefreshError(args=(str(), {"error": "invalid_grant"})):
+                return True
+            case _:
+                return "Reauthentication is needed" in str(error)
+
+    @staticmethod
+    def _token_fetch_failed(error: Exception) -> bool:
+        from google.auth.exceptions import RefreshError, TransportError
+
+        return isinstance(error, (RefreshError, TransportError))
+
     def _acquire_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey) -> asyncio.Lock:
         """Increment the refcount and return the lock for ``credential_cache_key``.
 
@@ -506,6 +584,51 @@ class VertexBase:
             return cached_entry
         return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None)
 
+    def _load_auth_or_earlier_key_file(
+        self,
+        credentials: VERTEX_CREDENTIALS_TYPES | None,
+        project_id: str | None,
+        credential_cache_key: _VertexCredentialCacheKey,
+    ) -> tuple[_VertexCredentialsObject | None, str | None]:
+        try:
+            return self.load_auth(credentials=credentials, project_id=project_id)
+        except Exception as e:
+            earlier: Final = (
+                None
+                if self._token_fetch_failed(e)
+                else self._credentials_from_an_earlier_key_file(credential_cache_key, e)
+            )
+            if earlier is None:
+                verbose_logger.exception(
+                    "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: %s",
+                    e,
+                )
+                raise
+            return earlier
+
+    def _credentials_from_an_earlier_key_file(
+        self, credential_cache_key: _VertexCredentialCacheKey, error: Exception
+    ) -> tuple[_VertexCredentialsObject, str | None] | None:
+        with self._credentials_cache_lock:
+            earlier_key: Final = next(
+                (
+                    key
+                    for key in reversed(self._credentials_project_mapping)
+                    if _is_other_version_of_key_file(key, credential_cache_key)
+                ),
+                None,
+            )
+        if earlier_key is None:
+            return None
+        earlier_credentials, earlier_project_id = self._unpack_cached_credentials(earlier_key)
+        if earlier_credentials is None:
+            return None
+        verbose_logger.warning(
+            "Vertex: the replaced credentials file did not load, so the credentials loaded before it stay in use: %s",
+            error,
+        )
+        return earlier_credentials, earlier_project_id
+
     def _cache_credentials(
         self,
         credential_cache_key: _VertexCredentialCacheKey,
@@ -548,14 +671,9 @@ class VertexBase:
         credential_cache_key: _VertexCredentialCacheKey,
     ) -> tuple[_VertexCredentialsObject, str | None]:
         """Load credentials via load_auth (in thread) and cache the result."""
-        try:
-            _credentials, credential_project_id = await asyncify(self.load_auth)(
-                credentials=credentials,
-                project_id=project_id,
-            )
-        except Exception as e:
-            verbose_logger.exception("Failed to load vertex credentials: %s", str(e))
-            raise
+        _credentials, credential_project_id = await asyncify(self._load_auth_or_earlier_key_file)(
+            credentials, project_id, credential_cache_key
+        )
         if _credentials is None:
             raise ValueError("Could not resolve credentials")
         self._cache_credentials(credential_cache_key, _credentials, credential_project_id)
@@ -842,7 +960,8 @@ class VertexBase:
         Handle reauthentication when credentials refresh fails.
 
         This method clears the cached credentials and attempts to reload them once.
-        It should only be called when "Reauthentication is needed" error occurs.
+        It should only be called when Google rejects the cached credentials' refresh,
+        e.g. "Reauthentication is needed" or a key file rotated to a new key.
 
         Args:
             credentials: The original credentials
@@ -863,12 +982,13 @@ class VertexBase:
         # Clear the cached credentials
         with self._credentials_cache_lock:
             self._credentials_project_mapping.pop(credential_cache_key, None)
+            self._credentials_project_mapping.pop((credential_cache_key[0], project_id), None)
 
         # Retry once with _retry_reauth=True to prevent infinite recursion
         try:
             return self.get_access_token(
                 credentials=credentials,
-                project_id=project_id,
+                project_id=credential_cache_key[1],
                 _retry_reauth=True,
             )
         except Exception as retry_error:
@@ -897,6 +1017,7 @@ class VertexBase:
 
         with self._credentials_cache_lock:
             self._credentials_project_mapping.pop(credential_cache_key, None)
+            self._credentials_project_mapping.pop((credential_cache_key[0], project_id), None)
 
         try:
             (
@@ -909,14 +1030,14 @@ class VertexBase:
             )
             if project_id is None and isinstance(credential_project_id, str):
                 project_id = credential_project_id
-                cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials
-                resolved_cache_key: Final = (cache_credentials, project_id)
+                resolved_cache_key: Final = (credential_cache_key[0], project_id)
                 # Always overwrite — any pre-existing entry at the resolved key
                 # references the OLD credentials object we just replaced, and
                 # leaving it would force the next request to do a redundant
                 # refresh/reauth before realizing the cached creds are stale.
                 self._cache_credentials(resolved_cache_key, _credentials, credential_project_id)
 
+            await asyncify(self._refresh_if_still_needed)(_credentials, self._needs_blocking_refresh)
             if _credentials.token is None or not isinstance(_credentials.token, str):
                 raise ValueError(
                     f"Could not resolve credentials token. Got None or non-string token (type={type(_credentials.token).__name__})"
@@ -958,11 +1079,32 @@ class VertexBase:
             Tuple of (access_token, project_id)
         """
 
-        # Convert dict credentials to string for caching
-        cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials
+        cache_credentials: Final = _credential_cache_identity(credentials)
         credential_cache_key: Final = (cache_credentials, project_id)
         verbose_logger.debug("Checking cached credentials for project_id: %s", project_id)
+        if _retry_reauth:
+            return self._resolve_access_token(credentials, project_id, credential_cache_key, _retry_reauth=True)
+        flight, leads = self._join_sync_token_flight(credential_cache_key)
+        if not leads:
+            return flight.result()
+        try:
+            flight.set_result(self._resolve_access_token(credentials, project_id, credential_cache_key))
+        except BaseException as e:
+            flight.set_exception(e)
+            raise
+        finally:
+            with self._credentials_cache_lock:
+                self._sync_token_flights.pop(credential_cache_key, None)
+        return flight.result()
 
+    def _resolve_access_token(
+        self,
+        credentials: VERTEX_CREDENTIALS_TYPES | None,
+        project_id: str | None,
+        credential_cache_key: _VertexCredentialCacheKey,
+        _retry_reauth: bool = False,
+    ) -> tuple[str, str]:
+        cache_credentials: Final = credential_cache_key[0]
         _credentials, credential_project_id = self._unpack_cached_credentials(credential_cache_key)
         if _credentials is not None:
             verbose_logger.debug("Cached credentials found for project_id: %s.", project_id)
@@ -976,14 +1118,9 @@ class VertexBase:
                 "Credential cache key not found for project_id: %s, loading new credentials", project_id
             )
 
-            try:
-                _credentials, credential_project_id = self.load_auth(credentials=credentials, project_id=project_id)
-            except Exception as e:
-                verbose_logger.exception(
-                    "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: %s",
-                    e,
-                )
-                raise e
+            _credentials, credential_project_id = self._load_auth_or_earlier_key_file(
+                credentials, project_id, credential_cache_key
+            )
 
             if _credentials is None:
                 raise ValueError(
@@ -1005,23 +1142,19 @@ class VertexBase:
             raise ValueError("Credentials are None after loading")
 
         if _credentials.expired:
-            with self._sync_refresh_lock:
-                # Double-check after acquiring lock
-                if _credentials.expired:
-                    try:
-                        verbose_logger.debug("Credentials expired, refreshing")
-                        self.refresh_auth(_credentials)
-                    except Exception as e:
-                        # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login`
-                        # in this case, we should try to reload the credentials by clearing the cache and retrying
-                        if "Reauthentication is needed" in str(e) and not _retry_reauth:
-                            return self._handle_reauthentication(
-                                credentials=credentials,
-                                project_id=project_id,
-                                credential_cache_key=credential_cache_key,
-                                error=e,
-                            )
-                        raise e
+            try:
+                self._refresh_if_still_needed(_credentials, lambda loaded: loaded.expired)
+            except Exception as e:
+                # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login`
+                # or rotated the key file, so we should try to reload the credentials by clearing the cache and retrying
+                if self._refresh_rejected(e) and not _retry_reauth:
+                    return self._handle_reauthentication(
+                        credentials=credentials,
+                        project_id=project_id,
+                        credential_cache_key=credential_cache_key,
+                        error=e,
+                    )
+                raise e
 
         ## VALIDATION STEP
         if _credentials.token is None or not isinstance(_credentials.token, str):
@@ -1048,7 +1181,7 @@ class VertexBase:
         """
         from google.auth.credentials import TokenState
 
-        cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials
+        cache_credentials: Final = _credential_cache_identity(credentials)
         credential_cache_key: Final = (cache_credentials, project_id)
 
         # === FAST PATH (no lock) ===
@@ -1128,9 +1261,9 @@ class VertexBase:
                     # Token is expired or missing — must block until refresh completes.
                     try:
                         verbose_logger.debug("Credentials expired, refreshing")
-                        await asyncify(self.refresh_auth)(_credentials)
+                        await asyncify(self._refresh_if_still_needed)(_credentials, self._needs_blocking_refresh)
                     except Exception as e:
-                        if "Reauthentication is needed" in str(e):
+                        if self._refresh_rejected(e):
                             verbose_logger.debug("Reauthentication needed, clearing cache and retrying")
                             return await self._handle_reauthentication_async(
                                 credentials=credentials,

litellm/utils.py

@@ -11031,10 +11031,10 @@ class ProviderConfigManager:
     def _get_vertex_ai_vector_stores_config(api_type: str | None) -> BaseVectorStoreConfig | None:
         if api_type == "rag_api" or api_type is None:
             from litellm.llms.vertex_ai.vector_stores.rag_api.transformation import (
-                VertexVectorStoreConfig,
+                vertex_rag_api_vector_store_config,
             )
 
-            return VertexVectorStoreConfig()
+            return vertex_rag_api_vector_store_config
         if api_type == "search_api":
             from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
                 vertex_search_api_vector_store_config,

tests/integration/providers/test_vector_store_search_wire.py

@@ -4,7 +4,9 @@ from typing import Final
 from urllib.parse import urlsplit
 
 import httpx
+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 openai import OpenAI
 from pydantic import BaseModel, JsonValue, TypeAdapter
@@ -30,6 +32,9 @@ _ALIAS_BODY: Final[dict[str, JsonValue]] = {
     "rewrite_query": None,
 }
 _ALIAS_REQUEST_BODY: Final[dict[str, JsonValue]] = {"query": "c"}
+_VERTEX_PROJECT: Final = "scripted-search-project"
+_VERTEX_TOKEN: Final = "scripted-search-token"
+_VERTEX_SEARCH_COUNT: Final = 3
 _SEARCH_RESULT: Final[dict[str, JsonValue]] = {
     "file_id": "file-search-wire",
     "filename": "search.txt",
@@ -139,3 +144,61 @@ def test_vector_store_search_forwards_all_search_fields_on_both_routes(gateway:
             ("POST", f"/v1{search_path}"),
             ("POST", f"/v1{search_path}"),
         ]
+
+
+@pytest.mark.parametrize(
+    ("api_type", "search_target", "search_reply"),
+    (
+        (
+            "search_api",
+            "/store:search",
+            {"results": [{"id": "doc-1", "document": {"id": "doc-1", "derivedStructData": {"title": "hit"}}}]},
+        ),
+        ("rag_api", "/store:retrieveContexts", {"contexts": {"contexts": []}}),
+    ),
+)
+def test_vertex_vector_store_reuses_one_google_token_across_searches(
+    gateway: Gateway, api_type: str, search_target: str, search_reply: dict[str, JsonValue]
+) -> None:
+    store_id: Final = f"vs_vertex_{uuid.uuid4().hex}"
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        if request.target == "/_oauth/token":
+            return Reply(
+                body=json.dumps({"access_token": _VERTEX_TOKEN, "token_type": "Bearer", "expires_in": 3600}).encode()
+            )
+        assert request.target == search_target, request.target
+        assert request.headers["authorization"] == f"Bearer {_VERTEX_TOKEN}"
+        return Reply(body=json.dumps(search_reply).encode())
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        registered: Final = gateway.request(
+            "POST",
+            "/vector_store/new",
+            {
+                "vector_store_id": store_id,
+                "custom_llm_provider": f"vertex_ai/{api_type}",
+                "litellm_params": {
+                    "api_base": f"{wire.url}/store",
+                    "vertex_project": _VERTEX_PROJECT,
+                    "vertex_location": "global",
+                    "vertex_credentials": service_account_json(_VERTEX_PROJECT, wire.url),
+                },
+            },
+        )
+        assert registered.status_code == 200, registered.text
+        scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id})
+
+        responses: Final = tuple(
+            gateway.request("POST", f"/v1/vector_stores/{store_id}/search", {"query": "hit"})
+            for _ in range(_VERTEX_SEARCH_COUNT)
+        )
+        assert [response.status_code for response in responses] == [200] * _VERTEX_SEARCH_COUNT, [
+            response.text for response in responses
+        ]
+
+        assert [request.target for request in wire.drain()] == [
+            "/_oauth/token",
+            *(search_target,) * _VERTEX_SEARCH_COUNT,
+        ]

tests/unit/integrations/gcs_pubsub/test_pub_sub.py

@@ -93,22 +93,7 @@ async def test_construct_request_headers_project_id_from_env(monkeypatch):
 
 
 def _authorized_user_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, project_id: str) -> str:
-    credentials_path: Final = tmp_path / "google-auth.json"
-    credentials_path.write_text(
-        json.dumps(
-            {
-                "type": "authorized_user",
-                "client_id": "unit-client",
-                "client_secret": "unit-secret",
-                "refresh_token": "unit-refresh-token",
-                "token": "unit-access-token",
-                "token_uri": "https://oauth2.googleapis.com/token",
-                "quota_project_id": "unit-project",
-                "expiry": "2099-01-01T00:00:00Z",
-            }
-        )
-    )
-    credentials_path_string: Final = str(credentials_path)
+    credentials_path_string: Final = str(tmp_path / "google-auth.json")
     credentials: Final = _CachedVertexCredentials(token="unit-access-token", quota_project_id=project_id)
     vertex_chat_completion: Final = litellm.vertex_chat_completion
     monkeypatch.setitem(

tests/unit/llms/vertex_ai/test_vertex_llm_base.py

@@ -5,10 +5,14 @@ import logging
 import os
 import pathlib
 import re
+import time
 import traceback
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta
-from threading import Event
-from typing import Final, Literal
+from functools import partial
+from threading import Event, Semaphore
+from typing import Final, Literal, NoReturn
 from unittest.mock import MagicMock, call, patch
 
 import google.auth.exceptions
@@ -2595,14 +2599,286 @@ def test_get_access_token_reraises_malformed_service_account_credentials_without
     ]
 
 
-def _authorized_user_cache_credentials(identity: str) -> str:
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.asyncio
+async def test_expired_credentials_shared_by_an_inferred_and_an_explicit_project_refresh_once(is_async: bool) -> None:
+    vertex: Final = VertexBase()
+    now: Final = datetime(2025, 1, 1)
+    credentials: Final = _authorized_user_cache_credentials("shared")
+    with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth:
+        def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+            if len(oauth.calls) == 1:
+                time.sleep(0.3)
+            return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600})
+
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
+        assert vertex.get_access_token(credentials, None) == ("token-1", "shared")
+        clock.return_value = now + timedelta(seconds=7200)
+        fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+        tokens: Final = tuple(await asyncio.gather(*(fetch(credentials, project) for project in (None, "shared"))))
+
+        assert tokens == (("token-2", "shared"), ("token-2", "shared"))
+        assert len(oauth.calls) == 2
+
+
+def test_callers_that_join_a_rejected_google_token_fetch_share_its_error_and_the_next_call_retries() -> None:
+    vertex: Final = VertexBase()
+    credentials: Final = _authorized_user_cache_credentials("revoked")
+    leader_fetching: Final = Event()
+    followers_calling: Final = Semaphore(0)
+
+    def follow() -> tuple[str, str]:
+        followers_calling.release()
+        return vertex.get_access_token(credentials, "unit-project")
+
+    def revoked_token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        if not leader_fetching.is_set():
+            leader_fetching.set()
+            assert all(followers_calling.acquire(timeout=5) for _ in range(4))
+            time.sleep(0.2)
+        return 400, {}, json.dumps({"error": "invalid_grant", "error_description": "Token has been expired or revoked."})
+
+    with responses.RequestsMock() as oauth, ThreadPoolExecutor(max_workers=5) as pool:
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=revoked_token_response)
+        leader: Final = pool.submit(vertex.get_access_token, credentials, "unit-project")
+        assert leader_fetching.wait(timeout=5)
+        followers: Final = tuple(pool.submit(follow) for _ in range(4))
+        errors: Final = tuple(call.exception(timeout=10) for call in (leader, *followers))
+
+        assert len(oauth.calls) == 1
+        assert all(isinstance(error, google.auth.exceptions.RefreshError) for error in errors)
+        with pytest.raises(google.auth.exceptions.RefreshError, match="invalid_grant"):
+            vertex.get_access_token(credentials, "unit-project")
+        assert len(oauth.calls) == 2
+
+
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("project", ["unit-project", None], ids=["explicit", "inferred"])
+@pytest.mark.asyncio
+async def test_callers_reload_a_rotated_key_file_once_when_google_rejects_the_old_key(
+    tmp_path: pathlib.Path, is_async: bool, project: str | None
+) -> None:
+    vertex: Final = VertexBase()
+    now: Final = datetime(2025, 1, 1)
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project"))
+    old_key_revoked: Final = Event()
+
+    def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        form: Final = request.body.decode() if isinstance(request.body, bytes) else str(request.body)
+        is_old_key: Final = "old-key" in form
+        if is_old_key and old_key_revoked.is_set():
+            return 400, {}, json.dumps({"error": "invalid_grant", "error_description": "Token has been expired or revoked."})
+        return 200, {}, json.dumps({"access_token": "old-token" if is_old_key else "new-token", "expires_in": 3600})
+
+    fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+    with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth:
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
+        assert await fetch(str(key_file), project) == ("old-token", "unit-project")
+        _rewrite_keeping_its_version(
+            key_file, _authorized_user_cache_credentials("new-key", quota_project="unit-project")
+        )
+        old_key_revoked.set()
+        clock.return_value = now + timedelta(seconds=7200)
+        assert await fetch(str(key_file), project) == ("new-token", "unit-project")
+        assert await fetch(str(key_file), project) == ("new-token", "unit-project")
+        assert len(oauth.calls) == 3
+
+
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    ("outage_status", "outage_body"),
+    [
+        (503, json.dumps({"error": "server_error", "error_description": "Service unavailable."})),
+        (502, "<html>Bad Gateway</html>"),
+    ],
+    ids=["503", "502"],
+)
+@pytest.mark.asyncio
+async def test_a_google_outage_during_refresh_keeps_the_cached_key_instead_of_reloading_it(
+    tmp_path: pathlib.Path, is_async: bool, outage_status: int, outage_body: str
+) -> None:
+    vertex: Final = VertexBase()
+    now: Final = datetime(2025, 1, 1)
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_cache_credentials("old-key"))
+    google_down: Final = Event()
+
+    def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        if google_down.is_set():
+            return outage_status, {}, outage_body
+        return 200, {}, json.dumps({"access_token": "old-token", "expires_in": 3600})
+
+    fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+    with (
+        patch("google.auth._helpers.utcnow", return_value=now) as clock,
+        patch("google.auth._exponential_backoff.time"),
+        responses.RequestsMock() as oauth,
+    ):
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
+        assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project")
+        _rewrite_keeping_its_version(key_file, _authorized_user_cache_credentials("new-key"))
+        google_down.set()
+        clock.return_value = now + timedelta(seconds=7200)
+        with pytest.raises(google.auth.exceptions.RefreshError):
+            await fetch(str(key_file), "unit-project")
+        assert not any("new-key" in str(call.request.body) for call in oauth.calls)
+
+
+def _rewrite_keeping_its_version(key_file: pathlib.Path, contents: str) -> None:
+    before: Final = key_file.stat()
+    key_file.write_text(contents)
+    os.utime(key_file, ns=(before.st_atime_ns, before.st_mtime_ns))
+
+
+def _move_into_place(key_file: pathlib.Path, contents: str) -> None:
+    before: Final = key_file.stat()
+    staged: Final = key_file.with_name("staged-key.json")
+    staged.write_text(contents)
+    os.utime(staged, ns=(before.st_atime_ns, before.st_mtime_ns))
+    os.replace(staged, key_file)
+
+
+def _rewrite_later(key_file: pathlib.Path, contents: str) -> None:
+    before: Final = key_file.stat()
+    key_file.write_text(contents)
+    os.utime(key_file, ns=(before.st_atime_ns, before.st_mtime_ns + 1_000_000_000))
+
+
+def _rewrite_longer(key_file: pathlib.Path, contents: str) -> None:
+    _rewrite_keeping_its_version(key_file, contents.replace("new-key", "new-key-v2"))
+
+
+def _old_or_new_token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+    form: Final = request.body.decode() if isinstance(request.body, bytes) else str(request.body)
+    return 200, {}, json.dumps({"access_token": "old-token" if "old-key" in form else "new-token", "expires_in": 3600})
+
+
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("from_adc_env", [False, True], ids=["vertex_credentials", "GOOGLE_APPLICATION_CREDENTIALS"])
+@pytest.mark.parametrize(
+    "swap_key_file",
+    [_move_into_place, _rewrite_later, _rewrite_longer],
+    ids=["moved_into_place", "rewritten_later", "rewritten_longer"],
+)
+@pytest.mark.asyncio
+async def test_a_swapped_key_file_is_used_on_the_next_call_while_the_old_key_still_works(
+    tmp_path: pathlib.Path,
+    monkeypatch: pytest.MonkeyPatch,
+    is_async: bool,
+    from_adc_env: bool,
+    swap_key_file: Callable[[pathlib.Path, str], None],
+) -> None:
+    vertex: Final = VertexBase()
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project"))
+    if from_adc_env:
+        monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(key_file))
+    credentials: Final = None if from_adc_env else str(key_file)
+
+    fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+    with responses.RequestsMock() as oauth:
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=_old_or_new_token_response)
+        assert await fetch(credentials, "unit-project") == ("old-token", "unit-project")
+        swap_key_file(key_file, _authorized_user_cache_credentials("new-key", quota_project="unit-project"))
+        assert await fetch(credentials, "unit-project") == ("new-token", "unit-project")
+        assert await fetch(credentials, "unit-project") == ("new-token", "unit-project")
+        assert len(oauth.calls) == 2
+
+
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.asyncio
+async def test_a_key_file_replaced_by_one_that_does_not_load_keeps_the_earlier_key_until_it_is_fixed(
+    tmp_path: pathlib.Path, is_async: bool
+) -> None:
+    vertex: Final = VertexBase()
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project"))
+
+    fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+    with responses.RequestsMock() as oauth:
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=_old_or_new_token_response)
+        assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project")
+        _move_into_place(key_file, '{"type": "authorized_user", "client_id": "half-wri')
+        assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project")
+        _move_into_place(key_file, _authorized_user_cache_credentials("repaired-key", quota_project="unit-project"))
+        assert await fetch(str(key_file), "unit-project") == ("new-token", "unit-project")
+        assert len(oauth.calls) == 2
+
+
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.asyncio
+async def test_a_key_file_that_does_not_load_raises_once_google_rejects_the_earlier_key(
+    tmp_path: pathlib.Path, is_async: bool
+) -> None:
+    vertex: Final = VertexBase()
+    now: Final = datetime(2025, 1, 1)
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project"))
+    old_key_revoked: Final = Event()
+
+    def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        if old_key_revoked.is_set():
+            return 400, {}, json.dumps({"error": "invalid_grant", "error_description": "Token has been expired or revoked."})
+        return _old_or_new_token_response(request)
+
+    fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+    with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth:
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
+        assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project")
+        _move_into_place(key_file, '{"type": "authorized_user", "client_id": "half-wri')
+        old_key_revoked.set()
+        clock.return_value = now + timedelta(seconds=7200)
+        with pytest.raises(google.auth.exceptions.RefreshError, match="invalid_grant"):
+            await fetch(str(key_file), "unit-project")
+
+
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    "google_unavailable",
+    [
+        lambda: (503, {}, json.dumps({"error": "server_error", "error_description": "Service unavailable."})),
+        lambda: _raise(requests.exceptions.ConnectionError("oauth2.googleapis.com is unreachable")),
+    ],
+    ids=["503", "unreachable"],
+)
+@pytest.mark.asyncio
+async def test_a_swapped_key_file_whose_first_token_fetch_fails_is_tried_again_on_the_next_call(
+    tmp_path: pathlib.Path, is_async: bool, google_unavailable: Callable[[], tuple[int, dict[str, str], str]]
+) -> None:
+    vertex: Final = VertexBase()
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_cache_credentials("old-key", quota_project="unit-project"))
+    google_down: Final = Event()
+
+    def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        return google_unavailable() if google_down.is_set() else _old_or_new_token_response(request)
+
+    fetch: Final = vertex.get_access_token_async if is_async else partial(asyncio.to_thread, vertex.get_access_token)
+    with patch("google.auth._exponential_backoff.time"), responses.RequestsMock() as oauth:
+        oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response)
+        assert await fetch(str(key_file), "unit-project") == ("old-token", "unit-project")
+        _move_into_place(key_file, _authorized_user_cache_credentials("new-key", quota_project="unit-project"))
+        google_down.set()
+        with pytest.raises((google.auth.exceptions.RefreshError, google.auth.exceptions.TransportError)):
+            await fetch(str(key_file), "unit-project")
+        google_down.clear()
+        assert await fetch(str(key_file), "unit-project") == ("new-token", "unit-project")
+        assert await fetch(str(key_file), "unit-project") == ("new-token", "unit-project")
+
+
+def _raise(error: Exception) -> NoReturn:
+    raise error
+
+
+def _authorized_user_cache_credentials(identity: str, quota_project: str | None = None) -> str:
     return json.dumps(
         {
             "type": "authorized_user",
             "client_id": identity,
             "client_secret": "synthetic-secret",
             "refresh_token": identity,
-            "quota_project_id": identity,
+            "quota_project_id": quota_project or identity,
         }
     )
 

tests/unit/vector_stores/test_main.py

@@ -7,13 +7,24 @@ executor, and it must never leak into litellm_params/kwargs where logging would
 model_dump() it (the #19550 serialization trap).
 """
 
+from __future__ import annotations
+
+import asyncio
 import json
+import pathlib
+import threading
+import time
+import uuid
 from collections.abc import Mapping
 from typing import Final
 from unittest.mock import MagicMock, patch
 
 import httpx
 import pytest
+import requests
+import requests_mock as requests_mock_lib
+import responses
+import respx
 
 import litellm
 import litellm.vector_stores.main as vector_stores_main
@@ -384,3 +395,156 @@ async def test_retrieving_a_missing_openai_vector_store_raises_not_found(is_asyn
 
     assert caught.value.status_code == 404
     assert "No vector store found with id 'vs_missing'." in caught.value.message
+
+
+GOOGLE_TOKEN_URL: Final = "https://oauth2.googleapis.com/token"
+
+
+def _authorized_user_credentials(refresh_token: str) -> str:
+    return json.dumps(
+        {
+            "type": "authorized_user",
+            "client_id": "unit-client-id",
+            "client_secret": "unit-client-secret",
+            "refresh_token": refresh_token,
+        }
+    )
+
+
+def _search_rag_store(api_base: str, vertex_credentials: str):
+    return litellm.vector_stores.asearch(
+        query="what is LiteLLM?",
+        vector_store_id="unit-store",
+        custom_llm_provider="vertex_ai/rag_api",
+        api_base=api_base,
+        vertex_project="unit-project",
+        vertex_location="us-central1",
+        vertex_credentials=vertex_credentials,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("api_type", "search_method", "search_response"),
+    (
+        ("search_api", ":search", {"results": []}),
+        ("rag_api", ":retrieveContexts", {"contexts": {"contexts": []}}),
+    ),
+)
+async def test_concurrent_vertex_searches_with_the_same_credentials_fetch_one_google_token(
+    api_type: str,
+    search_method: str,
+    search_response: dict[str, object],
+    respx_mock: respx.MockRouter,
+    requests_mock: requests_mock_lib.Mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    credentials: Final = _authorized_user_credentials(f"unit-refresh-{uuid.uuid4().hex}")
+
+    def slow_token_response(
+        request: requests_mock_lib.Request, context: requests_mock_lib.Context
+    ) -> dict[str, object]:
+        time.sleep(0.2)
+        return {"access_token": "ya29.unit-token", "expires_in": 3600, "token_type": "Bearer"}
+
+    token_route: Final = requests_mock.post(GOOGLE_TOKEN_URL, json=slow_token_response)
+    search_route: Final = respx_mock.post(f"https://vertex.test/{api_type}{search_method}").mock(
+        return_value=httpx.Response(200, json=search_response)
+    )
+
+    await asyncio.gather(
+        *(
+            litellm.vector_stores.asearch(
+                query="what is LiteLLM?",
+                vector_store_id="unit-store",
+                custom_llm_provider=f"vertex_ai/{api_type}",
+                api_base=f"https://vertex.test/{api_type}",
+                vertex_project="unit-project",
+                vertex_location="us-central1",
+                vertex_credentials=credentials,
+            )
+            for _ in range(5)
+        )
+    )
+
+    assert search_route.call_count == 5
+    assert token_route.call_count == 1
+    assert search_route.calls.last.request.headers["authorization"] == "Bearer ya29.unit-token"
+
+
+@pytest.mark.asyncio
+async def test_a_slow_google_token_fetch_for_one_rag_store_does_not_hold_up_another(
+    respx_mock: respx.MockRouter,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    slow_fetch_started: Final = threading.Event()
+    other_store_searched: Final = threading.Event()
+
+    def access_token_for(form: str) -> str:
+        if "slow-store" not in form:
+            return "ya29.fast"
+        slow_fetch_started.set()
+        return f"ya29.slow-waited-{other_store_searched.wait(timeout=5)}"
+
+    def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]:
+        form: Final = request.body.decode() if isinstance(request.body, bytes) else str(request.body)
+        token: Final = {"access_token": access_token_for(form), "expires_in": 3600, "token_type": "Bearer"}
+        return 200, {}, json.dumps(token)
+
+    slow_route: Final = respx_mock.post("https://vertex.test/slow:retrieveContexts").mock(
+        return_value=httpx.Response(200, json={"contexts": {"contexts": []}})
+    )
+    respx_mock.post("https://vertex.test/fast:retrieveContexts").mock(
+        return_value=httpx.Response(200, json={"contexts": {"contexts": []}})
+    )
+
+    with responses.RequestsMock() as google:
+        google.add_callback(responses.POST, GOOGLE_TOKEN_URL, callback=token_response)
+        slow_search: Final = asyncio.ensure_future(
+            _search_rag_store("https://vertex.test/slow", _authorized_user_credentials(f"slow-store-{uuid.uuid4().hex}"))
+        )
+        assert await asyncio.to_thread(slow_fetch_started.wait, 5)
+        await _search_rag_store("https://vertex.test/fast", _authorized_user_credentials(f"fast-store-{uuid.uuid4().hex}"))
+        other_store_searched.set()
+        await slow_search
+
+    assert slow_route.calls.last.request.headers["authorization"] == "Bearer ya29.slow-waited-True"
+
+
+@pytest.mark.asyncio
+async def test_a_rag_store_reloads_a_rotated_key_file_when_google_rejects_the_old_key(
+    tmp_path: pathlib.Path,
+    respx_mock: respx.MockRouter,
+    requests_mock: requests_mock_lib.Mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    key_file: Final = tmp_path / "vertex-key.json"
+    key_file.write_text(_authorized_user_credentials(f"old-key-{uuid.uuid4().hex}"))
+    old_key_deleted: Final = threading.Event()
+
+    def token_response(
+        request: requests_mock_lib.Request, context: requests_mock_lib.Context
+    ) -> dict[str, object]:
+        is_old_key: Final = "old-key" in request.text
+        if is_old_key and old_key_deleted.is_set():
+            context.status_code = 400
+            return {"error": "invalid_grant", "error_description": "Token has been expired or revoked."}
+        return {"access_token": "ya29.old" if is_old_key else "ya29.new", "expires_in": 0, "token_type": "Bearer"}
+
+    requests_mock.post(GOOGLE_TOKEN_URL, json=token_response)
+    search_route: Final = respx_mock.post("https://vertex.test/rotated:retrieveContexts").mock(
+        return_value=httpx.Response(200, json={"contexts": {"contexts": []}})
+    )
+
+    await _search_rag_store("https://vertex.test/rotated", str(key_file))
+    key_file.write_text(_authorized_user_credentials(f"new-key-{uuid.uuid4().hex}"))
+    old_key_deleted.set()
+    await _search_rag_store("https://vertex.test/rotated", str(key_file))
+
+    assert [call.request.headers["authorization"] for call in search_route.calls] == [
+        "Bearer ya29.old",
+        "Bearer ya29.new",
+    ]