fix(proxy): restore deprecated key auth cache lookups

This commit is contained in:
Leo 2026-03-28 16:47:33 +08:00
parent b8f7d61400
commit c04a1dc1a2
2 changed files with 63 additions and 5 deletions

View file

@ -2400,10 +2400,25 @@ async def _lookup_deprecated_key(
# Check cache first
cached = _deprecated_key_cache.get(hashed_token)
cached = _deprecated_key_cache.get(hashed_token)
if cached is not None:
active_token_id, cache_expires_at_ts, revoke_at_ts = cached
if now_ts < cache_expires_at_ts and now_ts < revoke_at_ts:
if len(cached) == 3:
active_token_id, cache_expires_at_ts, revoke_at_ts = cached
elif len(cached) == 2:
# Backward compatibility for cache entries written before
# revoke_at_ts was added to the in-memory tuple.
active_token_id, cache_expires_at_ts = cached
revoke_at_ts = float("inf")
else:
_deprecated_key_cache.pop(hashed_token, None)
active_token_id = None
cache_expires_at_ts = 0.0
revoke_at_ts = 0.0
if (
active_token_id is not None
and now_ts < cache_expires_at_ts
and now_ts < revoke_at_ts
):
return active_token_id
else:
_deprecated_key_cache.pop(hashed_token, None)
@ -2414,12 +2429,17 @@ async def _lookup_deprecated_key(
"token": hashed_token,
"revoke_at": {"gt": now},
},
select={"active_token_id": True},
select={"active_token_id": True, "revoke_at": True},
)
if deprecated_row and deprecated_row.active_token_id:
if (
deprecated_row
and deprecated_row.active_token_id
and deprecated_row.revoke_at is not None
):
_deprecated_key_cache[hashed_token] = (
deprecated_row.active_token_id,
now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS,
deprecated_row.revoke_at.timestamp(),
)
return deprecated_row.active_token_id
# Only cache positive results; negative lookups are fast on indexed columns

View file

@ -2,6 +2,7 @@ import asyncio
import json
import os
import sys
from datetime import datetime, timedelta, timezone
from typing import Tuple
from unittest.mock import ANY, AsyncMock, MagicMock, patch
@ -903,6 +904,43 @@ def test_proxy_admin_jwt_auth_handles_no_team_object():
assert result.end_user_id is None
@pytest.mark.asyncio
async def test_lookup_deprecated_key_uses_cache_on_subsequent_requests():
"""
Regression test for deprecated key auth during grace period.
The in-memory cache entry must be readable on the second lookup; otherwise
rotated keys intermittently fail auth with token_not_found_in_db.
"""
from types import SimpleNamespace
from litellm.proxy.utils import _deprecated_key_cache, _lookup_deprecated_key
hashed_old_key = "old-key-hash"
hashed_new_key = "new-key-hash"
revoke_at = datetime.now(timezone.utc) + timedelta(minutes=10)
_deprecated_key_cache.clear()
mock_db = MagicMock()
mock_db.litellm_deprecatedverificationtoken.find_first = AsyncMock(
return_value=SimpleNamespace(
active_token_id=hashed_new_key,
revoke_at=revoke_at,
)
)
try:
first_lookup = await _lookup_deprecated_key(mock_db, hashed_old_key)
second_lookup = await _lookup_deprecated_key(mock_db, hashed_old_key)
finally:
_deprecated_key_cache.clear()
assert first_lookup == hashed_new_key
assert second_lookup == hashed_new_key
mock_db.litellm_deprecatedverificationtoken.find_first.assert_awaited_once()
class TestJWTOAuth2Coexistence:
"""
Test that JWT and OAuth2 auth can coexist on the same instance.