mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Fix stale Vertex token refresh fallback
This commit is contained in:
parent
1a377bdd28
commit
87803d541d
2 changed files with 41 additions and 3 deletions
|
|
@ -892,9 +892,9 @@ class VertexBase:
|
|||
raise ValueError("Could not resolve project_id")
|
||||
current_token = _credentials.token
|
||||
if current_token is None or not isinstance(current_token, str):
|
||||
# Token is malformed despite STALE state — fall through
|
||||
# to INVALID path which will block on a full refresh.
|
||||
pass
|
||||
# Token is malformed despite STALE state — block on a full
|
||||
# refresh using the same path as INVALID credentials.
|
||||
token_state = TokenState.INVALID
|
||||
else:
|
||||
# Schedule a single background refresh — skip if one is
|
||||
# already in flight for this credential key.
|
||||
|
|
|
|||
|
|
@ -1559,6 +1559,44 @@ class TestVertexBase:
|
|||
|
||||
assert mock_refresh.called, "Background refresh should have been triggered"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_malformed_token_blocks_on_refresh(self):
|
||||
"""Malformed STALE tokens should refresh instead of failing validation."""
|
||||
from google.auth.credentials import TokenState
|
||||
|
||||
vertex_base = VertexBase()
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.token = None
|
||||
mock_creds.token_state = TokenState.STALE
|
||||
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, project = await vertex_base._ensure_access_token_async(
|
||||
credentials=credentials,
|
||||
project_id="project-1",
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
assert mock_refresh.called
|
||||
assert token == "refreshed-token"
|
||||
assert project == "project-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fresh_token_skips_refresh(self):
|
||||
"""Credentials not marked expired by google-auth should not trigger refresh."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue