From 00b49408995b6863c0aadc9e69f638947fd1b6dc Mon Sep 17 00:00:00 2001 From: Gabi Davar Date: Sat, 10 Oct 2026 12:28:13 +0300 Subject: [PATCH 1/2] 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 665ab771726..e0e27f9d544 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -716,6 +716,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 77508096e4a..085bb55dc59 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -8,11 +8,13 @@ import asyncio import json 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 @@ -482,6 +484,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): @@ -490,6 +494,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 17a92b31284..1781d3d8226 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -2202,6 +2202,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 + _ENCODED_VERTEX_KEY: Final = base64.b64encode( json.dumps({"type": "service_account", "private_key_id": "kid-123", "private_key": "MIIEvQIBADA???>>>"}).encode() From d231d0f5b360f951106e7fba404a73ff1dfb0ce6 Mon Sep 17 00:00:00 2001 From: Gabi Davar Date: Fri, 16 Oct 2026 23:04:19 +0300 Subject: [PATCH 2/2] test(vertex_ai): drive the expiry margin tests from a frozen clock and Google's loader The two regression tests patched VertexBase.load_auth and read the real clock, and the background refresh assertion raced a 50ms sleep. They now patch google.oauth2's from_service_account_info with a credential whose token state follows one frozen clock, the module reads that clock through time.time so the test freezes it without touching LiteLLM code, and the background refresh is awaited from the tracked task instead of slept on --- litellm/llms/vertex_ai/vertex_llm_base.py | 7 +- .../llms/vertex_ai/test_vertex_llm_base.py | 136 +++++++++++------- 2 files changed, 89 insertions(+), 54 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 085bb55dc59..c1dccefca37 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -7,6 +7,7 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S import asyncio import json import threading +import time from collections.abc import Mapping from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol @@ -71,6 +72,10 @@ class _VertexCredentialsObject(Protocol): def refresh(self, request: object) -> None: ... +def _utcnow() -> datetime: + return datetime.fromtimestamp(time.time(), tz=timezone.utc) + + class VertexBase: def __init__(self) -> None: super().__init__() @@ -499,7 +504,7 @@ class VertexBase: expiry: Final = getattr(credentials, "expiry", None) if not isinstance(expiry, datetime): return False - now: Final = datetime.now(timezone.utc) + now: Final = _utcnow() 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 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 1781d3d8226..2baec4ea74e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -1,5 +1,6 @@ import asyncio import base64 +import datetime import json import logging import os @@ -2201,6 +2202,8 @@ class TestVertexBase: assert token == "cached-token" assert not mock_get_lock.called, "Fast path should not acquire lock" + 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): @@ -2209,45 +2212,30 @@ class TestVertexBase: 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 + clock = _FakeClock(datetime.datetime(2026, 1, 1, 12, 0, 0)) + creds = _FakeGoogleCredentials(clock) 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, + patch("google.oauth2.service_account.Credentials.from_service_account_info", return_value=creds), + patch("time.time", side_effect=clock.timestamp), ): - - 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 + first_token, _ = await vertex_base.ensure_access_token_async( + credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai" + ) + clock.now = creds.expiry - datetime.timedelta(seconds=VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS / 2) token, project = await vertex_base.ensure_access_token_async( - credentials=credentials, - project_id="project-1", - custom_llm_provider="vertex_ai", + 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 + assert first_token == "token-1" + assert token == "token-2" + assert project == "project-1" + assert creds.refresh_calls == 2 @pytest.mark.asyncio async def test_stale_token_outside_safety_margin_is_served_while_refreshing(self): @@ -2255,45 +2243,87 @@ class TestVertexBase: 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 + stale_window = google_auth_helpers.REFRESH_THRESHOLD.total_seconds() + if VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS >= stale_window: + pytest.skip("the configured margin covers the whole STALE window, there is no zero-latency path to test") + seconds_left = (VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS + stale_window) / 2 + + clock = _FakeClock(datetime.datetime(2026, 1, 1, 12, 0, 0)) + creds = _FakeGoogleCredentials(clock) 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, + patch("google.oauth2.service_account.Credentials.from_service_account_info", return_value=creds), + patch("time.time", side_effect=clock.timestamp), ): - - def mock_refresh_impl(creds): - creds.token = "refreshed-token" - creds.token_state = TokenState.FRESH - - mock_refresh.side_effect = mock_refresh_impl + await vertex_base.ensure_access_token_async( + credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai" + ) + clock.now = creds.expiry - datetime.timedelta(seconds=seconds_left) token, _ = await vertex_base.ensure_access_token_async( - credentials=credentials, - project_id="project-1", - custom_llm_provider="vertex_ai", + credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai" ) + assert token == "token-1" + await asyncio.gather(*vertex_base._background_refresh_tasks.values()) - assert token == "near-expiry-token" + assert creds.refresh_calls == 2 + assert creds.token == "token-2" - await asyncio.sleep(0.05) - assert mock_refresh.called +class _FakeClock: + """The one clock both the module (through time.time) and the fake credential read.""" + + def __init__(self, now: datetime.datetime): + self.now = now + + def timestamp(self) -> float: + return self.now.replace(tzinfo=datetime.timezone.utc).timestamp() + + +class _FakeGoogleCredentials: + """Stands in for google.oauth2 credentials: the frozen clock decides FRESH / STALE / INVALID + the way google-auth does, and refresh() mints the next token an hour out.""" + + def __init__(self, clock: _FakeClock): + import datetime + + self._clock = clock + self._refresh_threshold = datetime.timedelta(minutes=3, seconds=45) + self.token: str | None = None + self.expiry = clock.now + self.refresh_calls = 0 + self.project_id = "project-1" + self.quota_project_id = "project-1" + + @property + def token_state(self): + from google.auth.credentials import TokenState + + if self.token is None or self._clock.now >= self.expiry: + return TokenState.INVALID + if self._clock.now >= self.expiry - self._refresh_threshold: + return TokenState.STALE + return TokenState.FRESH + + @property + def expired(self) -> bool: + return self._clock.now >= self.expiry - self._refresh_threshold + + @property + def valid(self) -> bool: + return self.token is not None and not self.expired + + def refresh(self, request) -> None: + import datetime + + self.refresh_calls += 1 + self.token = f"token-{self.refresh_calls}" + self.expiry = self._clock.now + datetime.timedelta(hours=1) _ENCODED_VERTEX_KEY: Final = base64.b64encode(