restore an explicit no-match policy

2026-09-22 00:31:44 +00:00 · 2026-04-11 12:39:09 -07:00 · 2026-04-11 12:39:09 -07:00 · bef94c74cd
commit bef94c74cd
parent eabb6a31f1
4 changed files with 294 additions and 12 deletions
--- 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

--- 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()

--- 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
--- 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"