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(