From b461a67d97b1559669bb149389a2163d44d9bc6f Mon Sep 17 00:00:00 2001 From: Satvik Sawhney Date: Fri, 14 Aug 2026 10:56:43 +0530 Subject: [PATCH] fix(vertex_ai): surface the google-auth install hint on the async token path vertex_llm_base.py defines GOOGLE_IMPORT_ERROR_MESSAGE and already raised it at seven import sites, but four `from google.auth.credentials import TokenState` sites were left unguarded. All four sit only on the async token path, and get_access_token_async reaches one before any guarded helper, so an install without google-auth answered every vertex_ai request with a bare ModuleNotFoundError wrapped as a 500 while the sync path reported the actionable hint. Guarding only the site that fires today would close the reported bug, but the other three sit in front of it on the same path and would reopen it under any refactor that reaches a helper first. The new tests fail on unpatched code with a regex mismatch against the raw ModuleNotFoundError, and the parametrized case covers each helper separately so removing any single guard fails a named case. --- litellm/llms/vertex_ai/vertex_llm_base.py | 17 ++++++++--- .../llms/vertex_ai/test_vertex_llm_base.py | 30 +++++++++++++++++++ 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 445e34966a9..fd16427fabe 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -37,6 +37,15 @@ else: GoogleCredentialsObject = Any +def _import_token_state() -> "type[TokenState]": + try: + from google.auth.credentials import TokenState + except ImportError: + raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) + + return TokenState + + class VertexBase: def __init__(self) -> None: super().__init__() @@ -397,7 +406,7 @@ class VertexBase: Look up cached credentials and return (token, project_id) if the token is FRESH. Returns None if not cached or not fresh. """ - from google.auth.credentials import TokenState + TokenState: Final = _import_token_state() creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) if ( @@ -423,7 +432,7 @@ class VertexBase: credentials object so the caller can schedule a background refresh without holding the per-key async lock. """ - from google.auth.credentials import TokenState + TokenState: Final = _import_token_state() creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) if creds is None: @@ -457,7 +466,7 @@ class VertexBase: Falls back to expired/valid checks if token_state is unavailable (e.g. older google-auth versions or mock objects in tests). """ - from google.auth.credentials import TokenState as _TokenState + _TokenState: Final = _import_token_state() token_state: Final = getattr(credentials, "token_state", None) if isinstance(token_state, _TokenState): @@ -978,7 +987,7 @@ class VertexBase: only one coroutine refreshes while others wait on the lock. Uses native async refresh for service_account and authorized_user credentials. """ - from google.auth.credentials import TokenState + TokenState: Final = _import_token_state() cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key: Final = (cache_credentials, project_id) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 18fc239b7c6..75d9086ced1 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -2088,3 +2088,33 @@ class TestVertexBase: assert token == "cached-token" assert not mock_get_lock.called, "Fast path should not acquire lock" + + @pytest.mark.asyncio + async def test_missing_google_auth_raises_actionable_import_error(self): + """Test that a missing google-auth surfaces the install hint, not a raw ModuleNotFoundError""" + vertex_base = VertexBase() + + with patch.dict("sys.modules", {"google.auth.credentials": None}): + with pytest.raises(ImportError, match="Google Cloud SDK not found"): + await vertex_base._ensure_access_token_async( + credentials={"type": "service_account", "project_id": "project-1"}, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + @pytest.mark.parametrize( + "call_helper", + [ + lambda vb: vb._try_get_cached_token(("key", None), "project-1"), + lambda vb: vb._try_get_usable_cached_token(("key", None), "project-1"), + lambda vb: vb._get_token_state(MagicMock()), + ], + ids=["try_get_cached_token", "try_get_usable_cached_token", "get_token_state"], + ) + def test_token_state_helpers_raise_actionable_import_error(self, call_helper): + """Test that every TokenState call site surfaces the install hint""" + vertex_base = VertexBase() + + with patch.dict("sys.modules", {"google.auth.credentials": None}): + with pytest.raises(ImportError, match="Google Cloud SDK not found"): + call_helper(vertex_base)