Merge d231d0f5b3 into da80552897

litellm/constants.py

@@ -737,6 +737,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

@@ -7,12 +7,15 @@ 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
 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
@@ -69,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__()
@@ -482,6 +489,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 +499,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 = _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
+
     async def _load_and_cache_credentials(
         self,
         credentials: VERTEX_CREDENTIALS_TYPES | None,

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,128 @@ 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):
+        """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 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()
+        credentials = {"type": "service_account", "project_id": "project-1"}
+
+        with (
+            patch("google.oauth2.service_account.Credentials.from_service_account_info", return_value=creds),
+            patch("time.time", side_effect=clock.timestamp),
+        ):
+            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"
+            )
+
+        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):
+        """A STALE token with comfortable time left keeps the zero-latency path."""
+        import datetime
+
+        from google.auth import _helpers as google_auth_helpers
+
+        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()
+        credentials = {"type": "service_account", "project_id": "project-1"}
+
+        with (
+            patch("google.oauth2.service_account.Credentials.from_service_account_info", return_value=creds),
+            patch("time.time", side_effect=clock.timestamp),
+        ):
+            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"
+            )
+            assert token == "token-1"
+            await asyncio.gather(*vertex_base._background_refresh_tasks.values())
+
+        assert creds.refresh_calls == 2
+        assert creds.token == "token-2"
+
+
+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(