Merge pull request #25570 from BerriAI/litellm_jwt_mapping

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

@@ -4161,6 +4161,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.
@@ -4261,6 +4280,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.",
@@ -4268,6 +4296,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()
 

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

tests/proxy_unit_tests/test_jwt_key_mapping.py

@@ -425,3 +425,421 @@ 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_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():
+    """
+    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"
+
+
+@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: 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
+# ──────────────────────────────────────────────
+
+
+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"