mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(proxy): return deprecated-key lookup result directly in get_data combined view (#30327)
The grace-period branch assigned the recursive get_data result (a
finished LiteLLM_VerificationTokenView) back into the variable that the
combined-view dict normalization then subscripts, raising TypeError on
every request made with a rotated key inside its grace window; auth
surfaced that as a 401. Return the recursive result directly instead.
Regression test drives the full get_data flow: old hash misses the view,
deprecated table resolves to the active token, and the call must return
the view object
(cherry picked from commit 5047eaf7f0)
This commit is contained in:
parent
bf6f134c24
commit
bcfed3e75d
2 changed files with 44 additions and 2 deletions
|
|
@ -3637,7 +3637,10 @@ class PrismaClient:
|
|||
db=self.db, hashed_token=hashed_token
|
||||
)
|
||||
if active_token_id:
|
||||
response = await self.get_data(
|
||||
# The recursive call returns a finished
|
||||
# LiteLLM_VerificationTokenView; the dict
|
||||
# normalization below would crash subscripting it.
|
||||
deprecated_response = await self.get_data(
|
||||
token=active_token_id,
|
||||
table_name="combined_view",
|
||||
query_type="find_unique",
|
||||
|
|
@ -3645,10 +3648,11 @@ class PrismaClient:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
check_deprecated=False,
|
||||
)
|
||||
if response is not None:
|
||||
if deprecated_response is not None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Deprecated key used during grace period"
|
||||
)
|
||||
return deprecated_response
|
||||
|
||||
if response is not None:
|
||||
if response["team_models"] is None:
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from __future__ import annotations
|
|||
|
||||
import hashlib
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -22,6 +23,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationTokenView
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
|
|
@ -476,3 +478,39 @@ async def test_get_data_logs_and_raises_on_db_error(
|
|||
)
|
||||
with pytest.raises(RuntimeError, match="network split"):
|
||||
await prisma_client.get_data(token="sk-broken", table_name="key")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_data_combined_view_returns_view_for_deprecated_key(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
"""Grace-period rotation, full get_data flow: the old hash misses the
|
||||
combined view, the deprecated-key table resolves it to the active token,
|
||||
and get_data must return the recursive lookup's finished view instead of
|
||||
re-running dict normalization on it (which raised TypeError and turned
|
||||
every grace-period request into a 401)."""
|
||||
old_hash = "hashed-old-token-grace-e2e"
|
||||
active_hash = "hashed-active-token-grace-e2e"
|
||||
active_row = {
|
||||
"token": active_hash,
|
||||
"team_models": None,
|
||||
"team_blocked": None,
|
||||
"team_members_with_roles": None,
|
||||
"user_id": None,
|
||||
"expires": None,
|
||||
}
|
||||
prisma_client.db.query_first = AsyncMock(side_effect=[None, active_row])
|
||||
prisma_client.db.litellm_deprecatedverificationtoken = MagicMock()
|
||||
prisma_client.db.litellm_deprecatedverificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
active_token_id=active_hash,
|
||||
revoke_at=datetime.now(timezone.utc) + timedelta(hours=1),
|
||||
)
|
||||
)
|
||||
|
||||
response = await prisma_client.get_data(
|
||||
token=old_hash, table_name="combined_view", query_type="find_unique"
|
||||
)
|
||||
|
||||
assert isinstance(response, LiteLLM_VerificationTokenView)
|
||||
assert response.token == active_hash
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue