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:
yuneng-jiang 2026-06-12 17:44:04 -07:00 committed by Yuneng Jiang
parent bf6f134c24
commit bcfed3e75d
No known key found for this signature in database
2 changed files with 44 additions and 2 deletions

View file

@ -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:

View file

@ -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