diff --git a/docs/my-website/docs/proxy/jwt_key_mapping.md b/docs/my-website/docs/proxy/jwt_key_mapping.md index 452bf821016..8937436e6d8 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 96e221a9ac0..77bf4d197f8 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 61c618eeb18..43646a11c63 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 66e5b3839bc..1b9d72d1b50 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"