From 8f25942ecfa7fb9c1f24a3dea44533794eb14f4e Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 12 May 2026 17:16:37 -0700 Subject: [PATCH] Litellm key rotation bug (#27756) * fix(proxy): resolve cache handling issues in _lookup_deprecated_key - Updated the in-memory cache for deprecated key lookups to store a 3-tuple (active_token_id, cache_expires_at_ts, revoke_at_ts) instead of a 2-tuple, ensuring proper unpacking and backward compatibility. - Removed duplicate cache reads and added logic to handle legacy cache entries gracefully. - Enhanced unit tests to cover scenarios for cache hits, DB misses, and respect for revoke_at timestamps, ensuring robust handling of the grace-period key-rotation feature. * refactor(proxy): streamline cache handling in _lookup_deprecated_key - Simplified the cache retrieval logic by directly unpacking the 3-tuple cache entries, removing the need for backward compatibility checks for 2-tuple entries. - Updated unit tests to ensure that pre-warmed 3-tuple cache entries are served correctly without unnecessary database lookups. * chore(ci): add new unit test for deprecated key grace period - Included `test_deprecated_key_grace_period.py` in the CI workflow to enhance coverage for deprecated key handling scenarios. * fix(proxy): remove unnecessary check for revoke_at in _lookup_deprecated_key - Eliminated the redundant check for None on revoke_at, streamlining the logic for handling deprecated keys in the cache. This change enhances the efficiency of the key lookup process. * test(proxy): add end-to-end tests for deprecated key lookup behavior - Introduced a new test class `TestDeprecatedKeyLookupDbE2E` to validate the behavior of deprecated key lookups against a real Prisma-backed database. - The test ensures that old key hashes resolve correctly and that repeated lookups utilize the in-memory cache without errors. - Cleaned up the `_lookup_deprecated_key` function by removing an unnecessary check for `revoke_at`, enhancing the efficiency of the key lookup process. --- .github/workflows/test-unit-proxy-db.yml | 1 + litellm/proxy/utils.py | 12 +- .../test_deprecated_key_grace_period.py | 177 ++++++++++++++++++ .../common_utils/test_key_rotation_e2e.py | 89 +++++++++ 4 files changed, 273 insertions(+), 6 deletions(-) create mode 100644 tests/proxy_unit_tests/test_deprecated_key_grace_period.py diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 5a9688db9c..bb44b00194 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -100,6 +100,7 @@ jobs: test-path: >- tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py + tests/proxy_unit_tests/test_deprecated_key_grace_period.py workers: 4 dist: loadscope timeout: 15 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2998555db1..559d5c99b9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2406,7 +2406,8 @@ def jsonify_object(data: dict) -> dict: return db_data -# In-memory cache for deprecated key lookups: maps old_token_hash -> (active_token_id, expires_at_ts) +# In-memory cache for deprecated key lookups: +# maps old_token_hash -> (active_token_id, cache_expires_at_ts, revoke_at_ts). # Avoids a DB query on every auth request for non-deprecated keys. # Bounded to prevent memory leaks from accumulated rotations. _deprecated_key_cache: LimitedSizeOrderedDict = LimitedSizeOrderedDict(max_size=1000) @@ -2428,26 +2429,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: return active_token_id - else: - _deprecated_key_cache.pop(hashed_token, None) + _deprecated_key_cache.pop(hashed_token, None) try: deprecated_row = await db.litellm_deprecatedverificationtoken.find_first( where={ "token": hashed_token, "revoke_at": {"gt": now}, - }, - select={"active_token_id": True}, + } ) if deprecated_row and deprecated_row.active_token_id: + revoke_at = deprecated_row.revoke_at _deprecated_key_cache[hashed_token] = ( deprecated_row.active_token_id, now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, + revoke_at.timestamp(), ) return deprecated_row.active_token_id # Only cache positive results; negative lookups are fast on indexed columns diff --git a/tests/proxy_unit_tests/test_deprecated_key_grace_period.py b/tests/proxy_unit_tests/test_deprecated_key_grace_period.py new file mode 100644 index 0000000000..a91ecf95f3 --- /dev/null +++ b/tests/proxy_unit_tests/test_deprecated_key_grace_period.py @@ -0,0 +1,177 @@ +""" +Tests for the grace-period key-rotation feature (MLI-6358). + +Two bugs are confirmed in LiteLLM v1.83.7-stable (upstream BerriAI/litellm#27193). +Both live in _lookup_deprecated_key() (litellm/proxy/utils.py): + + Bug 1 — duplicate cache read (cosmetic, no functional impact on its own): + The cache is fetched twice in a row with no state change between the calls. + + Bug 2 — cache stores a 2-tuple but unpacks as a 3-tuple: + WRITE: _deprecated_key_cache[hash] = (active_token_id, cache_expires_at_ts) + READ: active_token_id, cache_expires_at_ts, revoke_at_ts = cached # ValueError! + The ValueError is NOT inside the try/except, so it propagates up through + PrismaClient.get_data() (which re-raises), killing the auth request. + +The local demo script confirmed +that all three requests with the old key returned HTTP 401 immediately after +rotation even though the grace-period window was still open. +""" + +from datetime import datetime, timedelta, timezone +from typing import Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest + + +# ── helpers ─────────────────────────────────────────────────────────────────── + +HASHED_TOKEN = "165efe575c98fe7e65d98cb2de71b68842049e286afd33a92d3491c340216880" +ACTIVE_TOKEN_HASH = "abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890ab" + + +def _make_db(active_token_id: Optional[str]) -> MagicMock: + """Prisma db mock whose deprecated-token find_first returns the given id.""" + row = MagicMock() + row.active_token_id = active_token_id + row.revoke_at = datetime.now(timezone.utc) + timedelta(minutes=5) + db = MagicMock() + db.litellm_deprecatedverificationtoken = MagicMock() + db.litellm_deprecatedverificationtoken.find_first = AsyncMock( + return_value=row if active_token_id else None + ) + return db + + +# ── Bug 1: first call (DB path) ─────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_db_miss_returns_none(): + """Token absent from deprecated table → returns None without error.""" + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=None) + + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + + assert result is None + db.litellm_deprecatedverificationtoken.find_first.assert_called_once() + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_db_hit_returns_active_token_id(): + """ + First call (cold cache): DB row exists within grace window → returns + active_token_id correctly. The DB path itself works; the bug is on the + second call when the result is read back from cache. + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + + assert result == ACTIVE_TOKEN_HASH + db.litellm_deprecatedverificationtoken.find_first.assert_called_once() + + +# ── Bug 2: second call (cache path) ────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_cache_hit_returns_on_second_call(): + """ + Regression guard: after first call warms the cache with a 3-tuple, + second call should return from cache without raising. + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + # First call: cold cache → DB hit → warms cache with 3-tuple → succeeds + r1 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r1 == ACTIVE_TOKEN_HASH, "First call (DB path) must succeed" + + # Second call: cache hit path should succeed without DB access + r2 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r2 == ACTIVE_TOKEN_HASH + + # DB is queried exactly once; the second call never reaches it + assert db.litellm_deprecatedverificationtoken.find_first.call_count == 1 + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_pre_warmed_cache_returns(): + """ + Pre-warmed 3-tuple cache entry should be served directly from cache. + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + now_ts = datetime.now(timezone.utc).timestamp() + _deprecated_key_cache[HASHED_TOKEN] = ( + ACTIVE_TOKEN_HASH, + now_ts + 60, + now_ts + 300, + ) + + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert result == ACTIVE_TOKEN_HASH + + db.litellm_deprecatedverificationtoken.find_first.assert_not_called() + + +# ── End-to-end reproduction of the demo ────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_grace_period_three_requests_mirrors_demo(): + """ + Reproduces Step 5 of the local demo script: + + Request 1 (cache miss — DB lookup) → succeeds + Request 2 (cache hit) → succeeds + Request 3 (cache hit) → succeeds + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + r1 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r1 == ACTIVE_TOKEN_HASH, "Request 1 (DB path) should succeed" + + r2 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + r3 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r2 == ACTIVE_TOKEN_HASH + assert r3 == ACTIVE_TOKEN_HASH + + # DB hit only once; requests 2 and 3 never reach it + assert db.litellm_deprecatedverificationtoken.find_first.call_count == 1 + + +@pytest.mark.asyncio +async def test_cache_hit_respects_revoke_at_timestamp(): + """Cache entries should not remain valid past revoke_at even if cache TTL is still live.""" + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + now_ts = datetime.now(timezone.utc).timestamp() + # cache_expires_at is in the future, but revoke_at is already past. + _deprecated_key_cache[HASHED_TOKEN] = ( + ACTIVE_TOKEN_HASH, + now_ts + 60, + now_ts - 1, + ) + + db = _make_db(active_token_id=None) + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert result is None + db.litellm_deprecatedverificationtoken.find_first.assert_called_once() diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py index f6ef02a86d..d6e1d22fdd 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py @@ -13,7 +13,9 @@ Covers the critical gaps: import os import sys from datetime import datetime, timedelta, timezone +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 import pytest @@ -24,6 +26,11 @@ from litellm.proxy._types import ( LiteLLM_VerificationToken, ) from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager +from litellm.proxy.utils import ( + PrismaClient, + _deprecated_key_cache, + _lookup_deprecated_key, +) class TestMultiPodKeyRotation: @@ -557,3 +564,85 @@ class TestKeyRotationInitialization: assert acquire_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME assert release_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME + + +class TestDeprecatedKeyLookupDbE2E: + """DB-backed integration tests for deprecated key lookup behavior.""" + + @pytest.mark.asyncio + async def test_deprecated_key_grace_period_cache_hit_path(self): + """ + End-to-end validation against a real Prisma-backed DB: + - old key hash resolves through LiteLLM_DeprecatedVerificationToken + - repeated lookups hit the in-memory deprecated-key cache + - no ValueError/401 regression on subsequent requests + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + pytest.skip("DATABASE_URL not set; skipping DB-backed key-rotation E2E test.") + db_url = cast(str, database_url) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + prisma_client = PrismaClient( + database_url=db_url, proxy_logging_obj=proxy_logging_obj + ) + + old_token_hash = f"old-{uuid4().hex}" + active_token_hash = f"active-{uuid4().hex}" + _deprecated_key_cache.clear() + + await prisma_client.connect() + try: + await prisma_client.db.litellm_verificationtoken.create( + data={ + "token": active_token_hash, + "models": [], + } + ) + + await prisma_client.db.litellm_deprecatedverificationtoken.create( + data={ + "token": old_token_hash, + "active_token_id": active_token_hash, + "revoke_at": datetime.now(timezone.utc) + timedelta(minutes=5), + } + ) + + # Request 1 (DB path) + Request 2/3 (cache-hit path) + r1 = await _lookup_deprecated_key( + db=prisma_client.db, + hashed_token=old_token_hash, + ) + r2 = await _lookup_deprecated_key( + db=prisma_client.db, + hashed_token=old_token_hash, + ) + r3 = await _lookup_deprecated_key( + db=prisma_client.db, + hashed_token=old_token_hash, + ) + + assert r1 == active_token_hash + assert r2 == active_token_hash + assert r3 == active_token_hash + + cached = _deprecated_key_cache.get(old_token_hash) + assert isinstance(cached, tuple) + assert len(cached) == 3 + finally: + # Best-effort cleanup for idempotent reruns. + try: + await prisma_client.db.litellm_deprecatedverificationtoken.delete_many( + where={"token": old_token_hash} + ) + except Exception: + pass + try: + await prisma_client.db.litellm_verificationtoken.delete_many( + where={"token": active_token_hash} + ) + except Exception: + pass + _deprecated_key_cache.clear() + await prisma_client.disconnect()