From f5fabe1fe4fd0f47a16acdd8f9b1d8b24d4f478b Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 11 Apr 2026 13:15:02 -0700 Subject: [PATCH] fix(jwt): cache __NO_MAPPING__ sentinel before raising 403 in REJECT mode REJECT mode was raising HTTPException immediately on a DB miss without writing the __NO_MAPPING__ sentinel, causing every subsequent rejected request to re-query the DB. Write the sentinel first so repeated rejections are served from cache within virtual_key_mapping_cache_ttl. Adds test asserting DB is not hit on the second reject after a cache-warm miss. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/auth/user_api_key_auth.py | 7 +++ .../proxy_unit_tests/test_jwt_key_mapping.py | 56 +++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index beb679ec316..8ff8ece60f1 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -673,6 +673,13 @@ async def _resolve_jwt_to_virtual_key( behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior if behavior == UnregisteredJWTClientBehavior.REJECT: + # Cache the miss before raising so repeated rejections are served from + # cache and don't re-query the DB on every request. + await user_api_key_cache.async_set_cache( + key=cache_key, + value="__NO_MAPPING__", + ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + ) raise HTTPException( status_code=403, detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 2456731889b..257b2e5020f 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -466,6 +466,62 @@ async def test_reject_behavior_raises_403_on_no_mapping(): assert "unknown@example.com" in exc_info.value.detail +@pytest.mark.asyncio +async def test_reject_behavior_caches_sentinel_after_db_miss(): + """ + On a fresh DB miss with REJECT, the __NO_MAPPING__ sentinel must be written + to cache so that subsequent rejected requests are served from cache and do + not re-query the DB. + """ + from litellm.proxy._types import UnregisteredJWTClientBehavior + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="email", + unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.REJECT, + virtual_key_mapping_cache_ttl=300, + ) + jwt_claims = {"email": "unknown@example.com"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None) + + user_api_key_cache = DualCache() + + with patch("litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock): + # First call — DB miss, should raise 403 and write sentinel + with pytest.raises(HTTPException) as exc_info: + await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert exc_info.value.status_code == 403 + + # Sentinel must now be in cache + cached = await user_api_key_cache.async_get_cache( + "jwt_key_mapping:email:unknown@example.com" + ) + assert cached == "__NO_MAPPING__" + + # Second call — must raise 403 from cache, no additional DB hit + prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock() + with pytest.raises(HTTPException) as exc_info2: + await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert exc_info2.value.status_code == 403 + prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() + + @pytest.mark.asyncio async def test_reject_behavior_raises_403_on_cached_no_mapping(): """