fix(jwt): implement unregistered_jwt_client_behavior for JWT→virtual key mapping

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
 

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,13 +4276,29 @@ 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.",
     )
     #########################################################
 
     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()
 

litellm/proxy/auth/user_api_key_auth.py

@@ -499,6 +499,95 @@ 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.
+
+    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,
+    )
+
+    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"]
+
+    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,
+        value=token_hash,
+        ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
+    )
+
+    verbose_proxy_logger.info(
+        "JWT Key Mapping (auto_register): created new virtual key for %s='%s'.",
+        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 +616,26 @@ 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.",
+            )
+        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(
@@ -537,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(
@@ -559,13 +669,51 @@ async def _resolve_jwt_to_virtual_key(
             parent_otel_span=parent_otel_span,
             proxy_logging_obj=proxy_logging_obj,
         )
-    else:
+
+    # 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:
+        # 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,
         )
-        return None
+        raise HTTPException(
+            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:
+        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),
+            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