From 18662b3c13e62a3b344a143668a5dcb7b947da34 Mon Sep 17 00:00:00 2001 From: yuneng Date: Tue, 22 Sep 2026 01:52:38 +0000 Subject: [PATCH] fix(jwt): exclude blocked keys from auto_register_map_existing_key reuse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 11 ++++++++--- .../proxy/auth/test_user_api_key_auth.py | 13 ++++++++++--- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 055bc2978af..452780d97ba 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -932,9 +932,14 @@ async def _latest_active_key_hash_for_user(prisma_client: PrismaClient, user_id: row: Final = await VerificationTokenRepository(prisma_client).table.find_first( where={ # mutable-ok: the prisma where clause contract is a plain dict "user_id": user_id, - "OR": [ # mutable-ok: prisma filter literal - {"expires": None}, # mutable-ok: prisma filter literal - {"expires": {"gt": datetime.now(timezone.utc)}}, # mutable-ok: prisma filter literal + "AND": [ # mutable-ok: prisma filter literal + {"OR": [{"blocked": False}, {"blocked": None}]}, # mutable-ok: prisma filter literal + { # mutable-ok: prisma filter literal + "OR": [ # mutable-ok: prisma filter literal + {"expires": None}, # mutable-ok: prisma filter literal + {"expires": {"gt": datetime.now(timezone.utc)}}, # mutable-ok: prisma filter literal + ] + }, ], }, order={"created_at": "desc"}, # mutable-ok: prisma order literal 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 f10e40cb608..ff1fb938536 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 @@ -2145,7 +2145,9 @@ def _auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler, **over @pytest.mark.asyncio async def test_auto_register_map_existing_key_reuses_users_latest_key(): """With auto_register_map_existing_key on, the mapping must point at the user's - most recently created non-expired key hash and nothing may be minted.""" + most recently created non-expired, non-blocked key hash and nothing may be minted.""" + from typing import Final + from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping prisma_client = MagicMock() @@ -2175,10 +2177,15 @@ async def test_auto_register_map_existing_key_reuses_users_latest_key(): find_first.assert_awaited_once() where = find_first.await_args.kwargs["where"] assert where["user_id"] == "validated-user" - assert {"expires": None} in where["OR"], f"non-expired keys must be included: {where}" - assert any(isinstance(entry.get("expires"), dict) and "gt" in entry["expires"] for entry in where["OR"]), ( + expires_or: Final = next(entry["OR"] for entry in where["AND"] if any("expires" in e for e in entry["OR"])) + assert {"expires": None} in expires_or, f"non-expired keys must be included: {where}" + assert any(isinstance(entry.get("expires"), dict) and "gt" in entry["expires"] for entry in expires_or), ( f"future-expiring keys must be included: {where}" ) + blocked_or: Final = next(entry["OR"] for entry in where["AND"] if any("blocked" in e for e in entry["OR"])) + assert {"blocked": False} in blocked_or and {"blocked": None} in blocked_or, ( + f"blocked keys must be excluded from reuse: {where}" + ) assert find_first.await_args.kwargs["order"] == {"created_at": "desc"} create_data = prisma_client.db.litellm_jwtkeymapping.create.await_args.kwargs["data"]