From c9a8846278916b021d610cd1598c6332438dda5a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 20 Oct 2026 12:01:53 -0700
Subject: [PATCH] fix(vertex_ai): reuse one Vertex search OAuth token across
 searches

Every vector store search built a new search config, so the token cache it carries was thrown away and each search fetched a new Google OAuth token with a blocking call on the event loop. Under a burst of searches that blocking fetch held back every HTTP/2 response. The provider config is now shared, so searches with the same credentials reuse one token until it expires.
---
 .../search_api/transformation.py              |  3 ++
 litellm/utils.py                              |  4 +--
 tests/integration/sdk/test_http2_wire.py      | 33 ++++++++++++++++++-
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py
index 7b058bd0334..d7676d54451 100644
--- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py
+++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py
@@ -461,3 +461,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
 
         input_cost_per_query: Final = model_info.get("input_cost_per_query") or 0.0
         return input_cost_per_query, 0.0
+
+
+vertex_search_api_vector_store_config: Final = VertexSearchAPIVectorStoreConfig()
diff --git a/litellm/utils.py b/litellm/utils.py
index 14ee35e571e..19c70b73e48 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -10120,10 +10120,10 @@ class ProviderConfigManager:
                 return VertexVectorStoreConfig()
             elif api_type == "search_api":
                 from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
-                    VertexSearchAPIVectorStoreConfig,
+                    vertex_search_api_vector_store_config,
                 )
 
-                return VertexSearchAPIVectorStoreConfig()
+                return vertex_search_api_vector_store_config
         elif litellm.LlmProviders.BEDROCK == provider:
             from litellm.llms.bedrock.vector_stores.transformation import (
                 BedrockVectorStoreConfig,
diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py
index 4ff38b82837..793dd8d490b 100644
--- a/tests/integration/sdk/test_http2_wire.py
+++ b/tests/integration/sdk/test_http2_wire.py
@@ -40,6 +40,7 @@ STREAM_CHUNKS: Final = 3
 HELD_SEARCHES: Final = 150
 HOLD_SECONDS: Final = 10.0
 _HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()})
+_TOKEN_REQUESTS: Final = queue.SimpleQueue[str]()
 VERTEX_SERVING_CONFIG_PATH: Final = (
     "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search"
 )
@@ -108,7 +109,8 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab
             send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]}
         )
         return
-    if scope["path"] == "/_oauth/token":
+    if scope["path"].endswith("/_oauth/token"):
+        _TOKEN_REQUESTS.put(scope["path"])
         await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"})
         return
     if scope["path"].endswith(":search"):
@@ -327,6 +329,35 @@ def test_default_async_client_stays_http1_after_vertex_search_used_http2(
     assert asyncio.run(_run()) == (["http/2"], "1.1")
 
 
+def test_vertex_search_fetches_one_token_for_many_searches(
+    monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
+) -> None:
+    import litellm
+
+    _trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
+    credentials: Final = service_account_json("p", f"{http2_tls_peer.url}/token-reuse")
+    search: Final = MappingProxyType(
+        {
+            "vector_store_id": "default_search",
+            "query": "ping",
+            "custom_llm_provider": "vertex_ai/search_api",
+            "api_base": f"{http2_tls_peer.url}{VERTEX_SERVING_CONFIG_PATH}",
+            "vertex_credentials": credentials,
+            "vertex_project": "p",
+        }
+    )
+
+    async def _searches() -> None:
+        await litellm.vector_stores.asearch(**search)
+        await asyncio.gather(*(litellm.vector_stores.asearch(**search) for _ in range(5)))
+
+    asyncio.run(_searches())
+    litellm.vector_stores.search(**search)
+
+    token_paths: Final = tuple(iter(lambda: None if _TOKEN_REQUESTS.empty() else _TOKEN_REQUESTS.get(), None))
+    assert token_paths.count("/token-reuse/_oauth/token") == 1
+
+
 def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca(
     monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
 ) -> None: