From bef94c74cdf79d479a72b440bdd680c92e66d7dd Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 11 Apr 2026 12:39:09 -0700 Subject: [PATCH 1/4] restore an explicit no-match policy --- docs/my-website/docs/proxy/jwt_key_mapping.md | 16 +- litellm/proxy/_types.py | 35 ++++ litellm/proxy/auth/user_api_key_auth.py | 99 ++++++++++- .../proxy_unit_tests/test_jwt_key_mapping.py | 156 ++++++++++++++++++ 4 files changed, 294 insertions(+), 12 deletions(-) diff --git a/docs/my-website/docs/proxy/jwt_key_mapping.md b/docs/my-website/docs/proxy/jwt_key_mapping.md index 452bf821016c..8937436e6d8b 100644 --- a/docs/my-website/docs/proxy/jwt_key_mapping.md +++ b/docs/my-website/docs/proxy/jwt_key_mapping.md @@ -58,26 +58,30 @@ Complete [OIDC JWT Auth setup](./token_auth.md) first — you need `JWT_PUBLIC_K ### Step 1. Configure the JWT claim to map on -Add `jwt_client_id_field` to your `litellm_jwtauth` config. This is the JWT claim LiteLLM uses as the lookup key: +Add `virtual_key_claim_field` to your `litellm_jwtauth` config. This is the JWT claim LiteLLM uses as the lookup key: ```yaml general_settings: master_key: sk-1234 enable_jwt_auth: True litellm_jwtauth: - team_id_jwt_field: "team_id" # existing team mapping (optional) + team_id_jwt_field: "team_id" # existing team mapping (optional) user_id_jwt_field: "sub" - jwt_client_id_field: "client_id" # 👈 claim used for key mapping - unregistered_jwt_client_behavior: "fallback_team_mapping" # see below + virtual_key_claim_field: "client_id" # 👈 claim used for key mapping + unregistered_jwt_client_behavior: "reject" # see below ``` +:::note Renamed field +The field was called `jwt_client_id_field` in earlier docs. Both names are accepted — `jwt_client_id_field` silently maps to `virtual_key_claim_field`. +::: + **`unregistered_jwt_client_behavior`** controls what happens when a JWT has no registered mapping: | Value | Behavior | |-------|----------| | `fallback_team_mapping` | Fall through to team-based JWT auth (default — backward compatible) | -| `reject` | Return 403 if no mapping found | -| `auto_register` | Auto-create a virtual key + mapping on first encounter | +| `reject` | Return 403 if no mapping found. Use this when every caller must be pre-registered. | +| `auto_register` | Auto-create a virtual key + mapping on first encounter. The new key has no model/budget restrictions; tighten it later with `/jwt_client/update`. | ### Step 2. Register a JWT client → virtual key mapping diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 96e221a9ac03..77bf4d197f82 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4157,6 +4157,25 @@ class JWTRoutingOverride(BaseModel): } +class UnregisteredJWTClientBehavior(str, enum.Enum): + """ + Controls what happens when `virtual_key_claim_field` is configured but the + JWT claim value has no registered mapping in `litellm_jwtkeymapping`. + + - fallback_team_mapping: Fall through to standard team-based JWT auth (default, + backward-compatible). + - reject: Immediately return HTTP 403. Use this when every valid JWT client + must have a pre-registered virtual key — unknown callers are denied. + - auto_register: Automatically create a new virtual key and mapping on first + encounter. The new key has no budget/model restrictions; admins can tighten + it later via /jwt_client/update. + """ + + FALLBACK_TEAM_MAPPING = "fallback_team_mapping" + REJECT = "reject" + AUTO_REGISTER = "auto_register" + + class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): """ A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth. @@ -4257,6 +4276,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=300, description="TTL (seconds) for caching JWT-to-virtual-key mapping lookups.", ) + unregistered_jwt_client_behavior: UnregisteredJWTClientBehavior = Field( + default=UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING, + description=( + "What to do when virtual_key_claim_field is set but the JWT claim value " + "has no registered mapping. 'fallback_team_mapping' (default): fall through " + "to team-based JWT auth. 'reject': return HTTP 403. " + "'auto_register': auto-create a virtual key and mapping on first encounter." + ), + ) routing_overrides: Optional[List[JWTRoutingOverride]] = Field( default=None, description="Optional claim-based routing overrides for JWT-shaped tokens. Matching rules route requests to oauth2 before default JWT flow.", @@ -4264,6 +4292,13 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): ######################################################### def __init__(self, **kwargs: Any) -> None: + # Backward-compat: jwt_client_id_field was renamed to virtual_key_claim_field + if "jwt_client_id_field" in kwargs: + if "virtual_key_claim_field" not in kwargs: + kwargs["virtual_key_claim_field"] = kwargs.pop("jwt_client_id_field") + else: + kwargs.pop("jwt_client_id_field") + # get the attribute names for this Pydantic model allowed_keys = LiteLLM_JWTAuth.__annotations__.keys() diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 61c618eeb18f..43646a11c63c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -499,6 +499,64 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( return api_key +async def _auto_register_jwt_mapping( + virtual_key_claim_field: str, + claim_value: str, + jwt_handler: JWTHandler, + prisma_client: PrismaClient, + user_api_key_cache: DualCache, + parent_otel_span: Optional[Span], + proxy_logging_obj: ProxyLogging, + cache_key: str, +) -> Optional[UserAPIKeyAuth]: + """ + Auto-register: create a new virtual key + mapping for an unrecognised JWT claim value. + The new key carries no model/budget restrictions; admins can tighten it later. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_helper_fn, + ) + + key_data = await generate_key_helper_fn( + request_type="key", + metadata={ + "auto_registered": True, + "jwt_claim_field": virtual_key_claim_field, + "jwt_claim_value": claim_value, + }, + ) + token_hash = key_data["token"] + + await prisma_client.db.litellm_jwtkeymapping.create( + data={ + "jwt_claim_name": virtual_key_claim_field, + "jwt_claim_value": claim_value, + "token": token_hash, + "created_by": "auto_register", + "updated_by": "auto_register", + } + ) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value=token_hash, + ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + ) + + verbose_proxy_logger.info( + f"JWT Key Mapping (auto_register): created new virtual key for " + f"{virtual_key_claim_field}='{claim_value}'." + ) + + return await get_key_object( + hashed_token=token_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _resolve_jwt_to_virtual_key( jwt_claims: dict, jwt_handler: JWTHandler, @@ -527,6 +585,12 @@ async def _resolve_jwt_to_virtual_key( cached_mapping = await user_api_key_cache.async_get_cache(cache_key) if cached_mapping == "__NO_MAPPING__": + behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior + if behavior == UnregisteredJWTClientBehavior.REJECT: + raise HTTPException( + status_code=403, + detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", + ) return None elif cached_mapping is not None: return await get_key_object( @@ -559,13 +623,36 @@ async def _resolve_jwt_to_virtual_key( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - else: - await user_api_key_cache.async_set_cache( - key=cache_key, - value="__NO_MAPPING__", - ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + + # No mapping found — apply no-match policy + behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior + + if behavior == UnregisteredJWTClientBehavior.REJECT: + raise HTTPException( + status_code=403, + detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", ) - return None + + if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER: + return await _auto_register_jwt_mapping( + virtual_key_claim_field=virtual_key_claim_field, + claim_value=str(claim_value), + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + cache_key=cache_key, + ) + + # FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the + # caller falls through to standard team-based JWT auth. + await user_api_key_cache.async_set_cache( + key=cache_key, + value="__NO_MAPPING__", + ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + ) + return None async def _user_api_key_auth_builder( # noqa: PLR0915 diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 66e5b3839bc0..1b9d72d1b501 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -425,3 +425,159 @@ async def test_create_success_returns_response_without_token(): assert isinstance(result, JWTKeyMappingResponse) assert "token" not in result.model_fields assert result.jwt_claim_name == "email" + + +# ────────────────────────────────────────────── +# Tests: unregistered_jwt_client_behavior +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_reject_behavior_raises_403_on_no_mapping(): + """ + When unregistered_jwt_client_behavior='reject' and no mapping exists, + _resolve_jwt_to_virtual_key must raise HTTP 403. + """ + 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, + ) + 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): + 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 + assert "unknown@example.com" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_reject_behavior_raises_403_on_cached_no_mapping(): + """ + When the negative-cache sentinel __NO_MAPPING__ is present and behavior is + 'reject', the function must also raise HTTP 403 (not return None silently). + """ + 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, + ) + jwt_claims = {"email": "unknown@example.com"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None) + + # Pre-populate the negative cache so the DB is not hit + user_api_key_cache = DualCache() + cache_key = "jwt_key_mapping:email:unknown@example.com" + await user_api_key_cache.async_set_cache(cache_key, "__NO_MAPPING__") + + with patch("litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock): + 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 + # DB must NOT have been hit (sentinel served from cache) + prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() + + +@pytest.mark.asyncio +async def test_auto_register_creates_key_and_mapping(): + """ + When unregistered_jwt_client_behavior='auto_register' and no mapping exists, + _resolve_jwt_to_virtual_key must create a key + mapping row and return a + UserAPIKeyAuth object. + """ + from litellm.proxy._types import UnregisteredJWTClientBehavior + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER, + virtual_key_mapping_cache_ttl=300, + ) + jwt_claims = {"sub": "new-user-42"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock() + + user_api_key_cache = DualCache() + mock_key_obj = UserAPIKeyAuth(token="hashed_auto_key", team_id=None) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ) as mock_get_key, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_gen_key: + mock_gen_key.return_value = {"token": "hashed_auto_key", "key": "sk-auto-key"} + mock_get_key.return_value = mock_key_obj + + result = 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 result == mock_key_obj + # Mapping row must have been created + prisma_client.db.litellm_jwtkeymapping.create.assert_called_once() + call_data = prisma_client.db.litellm_jwtkeymapping.create.call_args[1]["data"] + assert call_data["jwt_claim_name"] == "sub" + assert call_data["jwt_claim_value"] == "new-user-42" + assert call_data["token"] == "hashed_auto_key" + # Cache must now hold the token hash + cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:new-user-42") + assert cached == "hashed_auto_key" + + +# ────────────────────────────────────────────── +# Tests: backward-compat alias jwt_client_id_field +# ────────────────────────────────────────────── + + +def test_jwt_client_id_field_alias_maps_to_virtual_key_claim_field(): + """ + jwt_client_id_field (old doc name) must silently alias to virtual_key_claim_field. + """ + auth = LiteLLM_JWTAuth(jwt_client_id_field="azp") + assert auth.virtual_key_claim_field == "azp" + + +def test_jwt_client_id_field_does_not_raise_on_duplicate(): + """ + If both jwt_client_id_field and virtual_key_claim_field are supplied, + virtual_key_claim_field takes precedence and no error is raised. + """ + auth = LiteLLM_JWTAuth( + jwt_client_id_field="old_field", + virtual_key_claim_field="new_field", + ) + assert auth.virtual_key_claim_field == "new_field" From 5daf7494c7497cf4fe95614f263f22d804e7d77a Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 11 Apr 2026 12:53:35 -0700 Subject: [PATCH 2/4] fix(jwt): fix AUTO_REGISTER sentinel bypass, race condition, and inline import comment - AUTO_REGISTER now evicts stale __NO_MAPPING__ sentinel instead of silently returning None when cached under a prior fallback_team_mapping config - Race condition in _auto_register_jwt_mapping: catch P2002 unique-constraint violation on concurrent creates, fetch the winning mapping, proceed cleanly - Added comment on inline generate_key_helper_fn import explaining the circular dependency (key_management_endpoints imports user_api_key_auth at line 51) - 3 new tests: stale sentinel eviction, race condition winner fallback, and the existing auto_register happy path Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/auth/user_api_key_auth.py | 67 ++++++++-- .../proxy_unit_tests/test_jwt_key_mapping.py | 117 ++++++++++++++++++ 2 files changed, 173 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 43646a11c63c..beb679ec316c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -512,7 +512,14 @@ async def _auto_register_jwt_mapping( """ Auto-register: create a new virtual key + mapping for an unrecognised JWT claim value. The new key carries no model/budget restrictions; admins can tighten it later. + + Race safety: if two concurrent requests both reach here simultaneously (both saw + no mapping in the DB), one will win the unique-constraint race on + litellm_jwtkeymapping. The loser catches the conflict, fetches the winner's + mapping, and proceeds — no orphaned keys and no error surfaced to the caller. """ + # Inline import required: key_management_endpoints imports user_api_key_auth + # (line 51) so a module-level import here would create a circular dependency. from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) @@ -527,15 +534,38 @@ async def _auto_register_jwt_mapping( ) token_hash = key_data["token"] - await prisma_client.db.litellm_jwtkeymapping.create( - data={ - "jwt_claim_name": virtual_key_claim_field, - "jwt_claim_value": claim_value, - "token": token_hash, - "created_by": "auto_register", - "updated_by": "auto_register", - } - ) + try: + await prisma_client.db.litellm_jwtkeymapping.create( + data={ + "jwt_claim_name": virtual_key_claim_field, + "jwt_claim_value": claim_value, + "token": token_hash, + "created_by": "auto_register", + "updated_by": "auto_register", + } + ) + except Exception as e: + error_str = str(e).lower() + if "unique" in error_str or "p2002" in error_str: + # A concurrent request won the race — fetch the winning mapping and + # use its token. The key we just generated is orphaned but harmless; + # it will be excluded from spend tracking since nothing maps to it. + verbose_proxy_logger.debug( + "JWT Key Mapping (auto_register): unique conflict on create — " + "fetching winner's mapping for %s='%s'.", + virtual_key_claim_field, + claim_value, + ) + token_hash = await get_jwt_key_mapping_object( + jwt_claim_name=virtual_key_claim_field, + jwt_claim_value=claim_value, + prisma_client=prisma_client, + ) + if token_hash is None: + # Should not happen, but guard against a delete racing our fetch. + return None + else: + raise await user_api_key_cache.async_set_cache( key=cache_key, @@ -544,8 +574,9 @@ async def _auto_register_jwt_mapping( ) verbose_proxy_logger.info( - f"JWT Key Mapping (auto_register): created new virtual key for " - f"{virtual_key_claim_field}='{claim_value}'." + "JWT Key Mapping (auto_register): created new virtual key for %s='%s'.", + virtual_key_claim_field, + claim_value, ) return await get_key_object( @@ -591,6 +622,20 @@ async def _resolve_jwt_to_virtual_key( status_code=403, detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", ) + if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER and prisma_client is not None: + # Stale sentinel written under a prior fallback_team_mapping config — + # evict it and auto-register now that the policy has changed. + await user_api_key_cache.async_delete_cache(cache_key) + return await _auto_register_jwt_mapping( + virtual_key_claim_field=virtual_key_claim_field, + claim_value=str(claim_value), + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + cache_key=cache_key, + ) return None elif cached_mapping is not None: return await get_key_object( diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 1b9d72d1b501..2456731889b6 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -558,6 +558,123 @@ async def test_auto_register_creates_key_and_mapping(): assert cached == "hashed_auto_key" +@pytest.mark.asyncio +async def test_auto_register_triggers_on_stale_no_mapping_sentinel(): + """ + If the cache holds a stale __NO_MAPPING__ sentinel (written under a prior + fallback_team_mapping config) and behavior is now AUTO_REGISTER, the sentinel + must be evicted and auto-registration must run — not silently return None. + """ + 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.AUTO_REGISTER, + virtual_key_mapping_cache_ttl=300, + ) + jwt_claims = {"email": "alice@corp.com"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock() + + user_api_key_cache = DualCache() + # Seed the stale sentinel + await user_api_key_cache.async_set_cache( + "jwt_key_mapping:email:alice@corp.com", "__NO_MAPPING__" + ) + + mock_key_obj = UserAPIKeyAuth(token="hashed_auto_key", team_id=None) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ) as mock_get_key, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_gen_key: + mock_gen_key.return_value = {"token": "hashed_auto_key", "key": "sk-auto-key"} + mock_get_key.return_value = mock_key_obj + + result = 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, + ) + + # Must have auto-registered, not returned None + assert result == mock_key_obj + prisma_client.db.litellm_jwtkeymapping.create.assert_called_once() + + +@pytest.mark.asyncio +async def test_auto_register_race_condition_unique_conflict(): + """ + If two concurrent requests both call _auto_register_jwt_mapping and the + second hits a unique-constraint violation on create, it must fall back to + fetching the winner's mapping — no error surfaced to the caller. + """ + from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping + from litellm.proxy._types import UnregisteredJWTClientBehavior + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER, + virtual_key_mapping_cache_ttl=300, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock( + side_effect=Exception("Unique constraint failed (P2002)") + ) + # Simulate the winner's mapping already in DB after the conflict + winner_mapping = MagicMock() + winner_mapping.token = "winner_token_hash" + winner_mapping.is_active = True + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock( + return_value=winner_mapping + ) + + user_api_key_cache = DualCache() + mock_key_obj = UserAPIKeyAuth(token="winner_token_hash", team_id=None) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ) as mock_get_key, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "loser_token_hash", "key": "sk-loser"}, + ): + mock_get_key.return_value = mock_key_obj + + result = await _auto_register_jwt_mapping( + virtual_key_claim_field="sub", + claim_value="user-42", + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + cache_key="jwt_key_mapping:sub:user-42", + ) + + assert result == mock_key_obj + # Cache should hold the winner's token, not the loser's + cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42") + assert cached == "winner_token_hash" + mock_get_key.assert_called_once_with( + hashed_token="winner_token_hash", + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + # ────────────────────────────────────────────── # Tests: backward-compat alias jwt_client_id_field # ────────────────────────────────────────────── From f5fabe1fe4fd0f47a16acdd8f9b1d8b24d4f478b Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 11 Apr 2026 13:15:02 -0700 Subject: [PATCH 3/4] 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 beb679ec316c..8ff8ece60f13 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 2456731889b6..257b2e5020f2 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(): """ From 86d8147df58e744a38d54eeb8cb87f6d678879c9 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 11 Apr 2026 13:29:20 -0700 Subject: [PATCH 4/4] fix(jwt): enforce no-match policy when prisma_client is None MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The early `if prisma_client is None: return None` guard ran before the no-match policy check, silently bypassing REJECT and AUTO_REGISTER — every JWT client fell through to team auth regardless of configuration. Fix: treat prisma_client=None as a definitive DB miss and fall through to the same policy block as a real miss. REJECT now raises 403, AUTO_REGISTER raises 500 with a clear message (can't create keys without a DB), FALLBACK_TEAM_MAPPING returns None unchanged. Adds three tests: REJECT/403 with no DB, FALLBACK returns None with no DB, AUTO_REGISTER/500 with no DB. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/auth/user_api_key_auth.py | 27 ++++-- .../proxy_unit_tests/test_jwt_key_mapping.py | 89 +++++++++++++++++++ 2 files changed, 107 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 8ff8ece60f13..4d7363336766 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -646,14 +646,15 @@ async def _resolve_jwt_to_virtual_key( proxy_logging_obj=proxy_logging_obj, ) - if prisma_client is None: - return None - - token_hash = await get_jwt_key_mapping_object( - jwt_claim_name=virtual_key_claim_field, - jwt_claim_value=str(claim_value), - prisma_client=prisma_client, - ) + # Resolve the mapping from DB, or treat prisma_client=None as a definitive + # miss (no DB → no mapping can exist → apply no-match policy below). + token_hash: Optional[str] = None + if prisma_client is not None: + token_hash = await get_jwt_key_mapping_object( + jwt_claim_name=virtual_key_claim_field, + jwt_claim_value=str(claim_value), + prisma_client=prisma_client, + ) if token_hash is not None: await user_api_key_cache.async_set_cache( @@ -669,7 +670,7 @@ async def _resolve_jwt_to_virtual_key( proxy_logging_obj=proxy_logging_obj, ) - # No mapping found — apply no-match policy + # No mapping found (DB miss or no DB) — apply no-match policy. behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior if behavior == UnregisteredJWTClientBehavior.REJECT: @@ -686,6 +687,14 @@ async def _resolve_jwt_to_virtual_key( ) if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER: + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=( + "JWT Key Mapping: AUTO_REGISTER requires a database connection. " + "Configure a database or change unregistered_jwt_client_behavior." + ), + ) return await _auto_register_jwt_mapping( virtual_key_claim_field=virtual_key_claim_field, claim_value=str(claim_value), diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 257b2e5020f2..15792918aa8d 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -731,6 +731,95 @@ async def test_auto_register_race_condition_unique_conflict(): ) +# ────────────────────────────────────────────── +# Tests: prisma_client=None does not bypass no-match policy +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_reject_behavior_enforced_when_prisma_client_is_none(): + """ + When prisma_client is None and behavior is REJECT, a 403 must be raised — + not silently fallen through to team auth. + """ + 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, + ) + jwt_claims = {"email": "unknown@example.com"} + + user_api_key_cache = DualCache() + + with pytest.raises(HTTPException) as exc_info: + await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=None, # no DB + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert exc_info.value.status_code == 403 + assert "unknown@example.com" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_fallback_team_mapping_returns_none_when_prisma_client_is_none(): + """ + When prisma_client is None and behavior is FALLBACK_TEAM_MAPPING, the + function must return None (fall through to team auth) — not raise. + """ + 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.FALLBACK_TEAM_MAPPING, + ) + jwt_claims = {"email": "anyone@example.com"} + + result = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_auto_register_raises_500_when_prisma_client_is_none(): + """ + AUTO_REGISTER without a DB connection must raise HTTP 500 with a clear + message — it cannot create keys without a database. + """ + from litellm.proxy._types import UnregisteredJWTClientBehavior + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER, + ) + jwt_claims = {"sub": "new-user-42"} + + with pytest.raises(HTTPException) as exc_info: + await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert exc_info.value.status_code == 500 + assert "AUTO_REGISTER requires a database" in exc_info.value.detail + + # ────────────────────────────────────────────── # Tests: backward-compat alias jwt_client_id_field # ──────────────────────────────────────────────