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

@@ -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,

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,

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()