From 63c2e002cc80256ef0863520ff848b71adb0ce64 Mon Sep 17 00:00:00 2001 From: Gabi Davar Date: Sat, 10 Oct 2026 12:28:13 +0300 Subject: [PATCH] fix(vertex_ai): refresh a STALE token before it reaches its expiry edge VertexBase served a cached token while google-auth reported it STALE and only refreshed it in the background. STALE lasts until the exact second of expiry, so a request in the last second went out with a token Google already rejects with 401 ACCESS_TOKEN_EXPIRED. A proxy whose only Vertex traffic is a 15 minute cron hits this every few hours, because no request lands inside the 3m45s STALE window, nothing refreshes early, and the tick one hour after the last refresh lands on the edge Treat a STALE token with VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS or less left as INVALID so the caller blocks on a refresh. Tokens with more time left keep the zero latency path and the background refresh --- litellm/constants.py | 1 + litellm/llms/vertex_ai/vertex_llm_base.py | 13 +++ .../llms/vertex_ai/test_vertex_llm_base.py | 93 +++++++++++++++++++ 3 files changed, 107 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index 38d16d61d7b..2f5d681b778 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -707,6 +707,7 @@ ANTHROPIC_IMAGE_MAX_PIXELS: Final = 1_150_000 ANTHROPIC_IMAGE_PIXELS_PER_TOKEN: Final = 750 PDF_DATA_URL_PREFIX: Final = "data:application/pdf;base64," VERTEX_BATCH_PREDICTION_JOBS_ROUTE: Final = "batchPredictionJobs" +VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS: Final = get_env_int("VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS", 60) ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { "low": 1, "medium": 5, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 1d5cbe09ba0..27c0cf4180c 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -9,11 +9,13 @@ import json import os import threading from collections.abc import Mapping +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse import litellm from litellm._logging import verbose_logger +from litellm.constants import VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.secret_managers.main import get_secret_str @@ -491,6 +493,8 @@ class VertexBase: token_state: Final = getattr(credentials, "token_state", None) if isinstance(token_state, _TokenState): + if token_state == _TokenState.STALE and self._expires_within_safety_margin(credentials): + return _TokenState.INVALID return token_state # Fallback for credentials without a real token_state (e.g. mocks) if getattr(credentials, "expired", True): @@ -499,6 +503,15 @@ class VertexBase: return _TokenState.FRESH return _TokenState.INVALID + @staticmethod + def _expires_within_safety_margin(credentials: _VertexCredentialsObject) -> bool: + expiry: Final = getattr(credentials, "expiry", None) + if not isinstance(expiry, datetime): + return False + now: Final = datetime.now(timezone.utc) + remaining: Final = expiry - (now if expiry.tzinfo is not None else now.replace(tzinfo=None)) + return remaining.total_seconds() <= VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS + async def _load_and_cache_credentials( self, credentials: VERTEX_CREDENTIALS_TYPES | None, diff --git a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py index 514728137cf..e51dd854a9d 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -2184,6 +2184,99 @@ class TestVertexBase: assert token == "cached-token" assert not mock_get_lock.called, "Fast path should not acquire lock" + @pytest.mark.asyncio + async def test_stale_token_inside_safety_margin_refreshes_before_use(self): + """A token a few seconds from expiry must be refreshed before it is handed out. + Google rejects it with ACCESS_TOKEN_EXPIRED at the edge, so serving it with only + a background refresh fails the request that triggered the refresh.""" + import datetime + + from google.auth import _helpers as google_auth_helpers + from google.auth.credentials import TokenState + + from litellm.constants import VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "about-to-expire-token" + mock_creds.token_state = TokenState.STALE + mock_creds.expiry = google_auth_helpers.utcnow() + datetime.timedelta( + seconds=VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS / 2 + ) + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object(vertex_base, "load_auth", return_value=(mock_creds, "project-1")), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + creds.expiry = google_auth_helpers.utcnow() + datetime.timedelta(hours=1) + + mock_refresh.side_effect = mock_refresh_impl + + token, project = await vertex_base.ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert token == "refreshed-token" + assert project == "project-1" + assert mock_refresh.call_count == 1 + + @pytest.mark.asyncio + async def test_stale_token_outside_safety_margin_is_served_while_refreshing(self): + """A STALE token with comfortable time left keeps the zero-latency path.""" + import datetime + + from google.auth import _helpers as google_auth_helpers + from google.auth.credentials import TokenState + + from litellm.constants import VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "near-expiry-token" + mock_creds.token_state = TokenState.STALE + mock_creds.expiry = google_auth_helpers.utcnow() + datetime.timedelta( + seconds=VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS * 2 + ) + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object(vertex_base, "load_auth", return_value=(mock_creds, "project-1")), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + token, _ = await vertex_base.ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert token == "near-expiry-token" + + await asyncio.sleep(0.05) + + assert mock_refresh.called + def test_get_access_token_reraises_malformed_service_account_credentials_without_caching(caplog): vertex_base: Final = VertexBase()