From c04a1dc1a2ce78d190e832f3b9343fcd73e613db Mon Sep 17 00:00:00 2001
From: Leo <leo48io@gmail.com>
Date: Sat, 28 Mar 2026 16:47:33 +0800
Subject: [PATCH] fix(proxy): restore deprecated key auth cache lookups

---
 litellm/proxy/utils.py                        | 30 ++++++++++++---
 .../proxy/auth/test_user_api_key_auth.py      | 38 +++++++++++++++++++
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index e15b48577de..ffa0e89bf21 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -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
diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
index ec7f3fc480c..ddeec2ff578 100644
--- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
+++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py
@@ -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.