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

[shivamrawat1]

Relevant issues

unregistered_jwt_client_behavior (reject / fallback_team_mapping / auto_register) was documented but never implemented.
_resolve_jwt_to_virtual_key had a single unconditional return None on no-match, so there was no way to enforce access control on
unmapped JWT clients — they always fell through to team auth. Separately, the doc still referenced the renamed field
jwt_client_id_field which LiteLLM_JWTAuth.init now rejects with a ValueError.

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Screenshots / Proof of Fix

Type

🐛 Bug Fix
📖 Documentation
✅ Test

Changes

• Added UnregisteredJWTClientBehavior enum and unregistered_jwt_client_behavior field to LiteLLM_JWTAuth (default:
  fallback_team_mapping — fully backward-compatible)
• reject: raises HTTP 403 on both fresh DB miss and cached NO_MAPPING sentinel
• auto_register: creates a new virtual key + mapping row on first encounter, proceeds as mapped
• Added backward-compat alias: jwt_client_id_field silently maps to virtual_key_claim_field so existing configs don't break
• Fixed doc to use virtual_key_claim_field with a deprecation note
• 5 new unit tests covering all three branches + the alias

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Apr 11, 2026 8:31pm

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing litellm_jwt_mapping (86d8147) with main (da75012)}

[greptile-apps]

Greptile Summary

This PR implements the previously documented but unimplemented unregistered_jwt_client_behavior field on LiteLLM_JWTAuth, adding REJECT, AUTO_REGISTER, and FALLBACK_TEAM_MAPPING dispatch logic to _resolve_jwt_to_virtual_key, along with a backward-compat alias for the renamed jwt_client_id_field. The stale-sentinel + AUTO_REGISTER path (a previously flagged gap) is now correctly handled by evicting the sentinel and re-running auto-registration.

Confidence Score: 4/5

Safe to merge with minor cleanup; remaining findings are style/P2 only.

The primary P1 gaps from the previous review round (AUTO_REGISTER silently ignoring a stale sentinel, and missing behavior enforcement) are now fixed. All remaining findings are P2: a direct Prisma write that bypasses the helper convention, and a non-optional type annotation that should be Optional. Neither blocks correct operation.

litellm/proxy/auth/user_api_key_auth.py — the new _auto_register_jwt_mapping helper contains the direct DB write and the type annotation gap.

Important Files Changed

Filename	Overview
litellm/proxy/auth/user_api_key_auth.py	Adds _auto_register_jwt_mapping helper and updates _resolve_jwt_to_virtual_key to dispatch to REJECT/AUTO_REGISTER/FALLBACK_TEAM_MAPPING; contains a direct Prisma write in the auth critical path and a type annotation inconsistency.
litellm/proxy/_types.py	Adds UnregisteredJWTClientBehavior enum and wires it into LiteLLM_JWTAuth; backward-compat alias for jwt_client_id_field is correctly handled in init before the invalid-key check.
tests/proxy_unit_tests/test_jwt_key_mapping.py	Adds thorough unit tests for all three behaviors, the stale sentinel + AUTO_REGISTER path, race-condition fallback, and prisma_client=None edge cases; all tests use mocks only (no real network calls).
docs/my-website/docs/proxy/jwt_key_mapping.md	Updates doc to use virtual_key_claim_field, adds deprecation note for jwt_client_id_field alias, and expands unregistered_jwt_client_behavior option descriptions; example config shows reject which differs from the code default (fallback_team_mapping) but is acceptable as an illustrative config.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[JWT auth request] --> B{virtual_key_claim_field set?}
    B -- No --> Z[return None → team auth]
    B -- Yes --> C{claim value in JWT?}
    C -- No --> Z
    C -- Yes --> D{Cache hit?}
    D -- token hash --> E[get_key_object → return UserAPIKeyAuth]
    D -- __NO_MAPPING__ sentinel --> F{unregistered_jwt_client_behavior}
    F -- REJECT --> G[raise HTTP 403]
    F -- AUTO_REGISTER + prisma ≠ None --> H[evict sentinel → _auto_register_jwt_mapping]
    F -- FALLBACK_TEAM_MAPPING or no prisma --> Z
    D -- miss --> I{DB lookup via get_jwt_key_mapping_object}
    I -- found --> J[cache token hash → get_key_object]
    I -- not found --> K{unregistered_jwt_client_behavior}
    K -- REJECT --> L[cache sentinel → raise HTTP 403]
    K -- AUTO_REGISTER + prisma ≠ None --> M[generate_key_helper_fn + jwtkeymapping.create]
    K -- AUTO_REGISTER + prisma = None --> N[raise HTTP 500]
    K -- FALLBACK_TEAM_MAPPING --> O[cache sentinel → return None]
    M -- success --> P[cache token hash → return UserAPIKeyAuth]
    M -- unique conflict --> Q[fetch winner mapping → cache + return]
    H --> M

Loading

_{Reviews (4): Last reviewed commit: "fix(jwt): enforce no-match policy when p..." | Re-trigger Greptile}

[greptile-apps]

AUTO_REGISTER silently falls through when cached __NO_MAPPING__ is present

When behavior == AUTO_REGISTER and the cache contains the __NO_MAPPING__ sentinel (written by a prior run under FALLBACK_TEAM_MAPPING), the function returns None — falling through to standard team-based JWT auth — instead of triggering auto-registration. Only REJECT is checked; AUTO_REGISTER is missed entirely in this branch.

if cached_mapping == "__NO_MAPPING__":
    behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
    if behavior == UnregisteredJWTClientBehavior.REJECT:
        raise HTTPException(status_code=403, ...)
    # AUTO_REGISTER lands here → returns None instead of registering
    return None

The fix is to also handle AUTO_REGISTER here (e.g., delete the stale sentinel and fall through to the DB path, or call _auto_register_jwt_mapping directly). Without this, switching a deployment from the default fallback_team_mapping to auto_register silently misbehaves until the cache TTL expires for every previously-seen client.

[greptile-apps]

Race condition creates orphaned keys under concurrent requests

Two simultaneous requests for the same previously-unseen JWT claim both pass the DB lookup (both see no mapping), both call _auto_register_jwt_mapping, and both reach generate_key_helper_fn. Because generate_key_helper_fn runs first and is not guarded, two virtual keys are created before either litellm_jwtkeymapping.create runs. The second create then fails on the @@unique([jwt_claim_name, jwt_claim_value]) constraint, leaving one key permanently orphaned in litellm_verificationtoken with no mapping row.

The fix is to catch the unique-constraint violation in _auto_register_jwt_mapping and, on conflict, fall back to fetching the row the winning request already inserted:

try:
    await prisma_client.db.litellm_jwtkeymapping.create(data={...})
except Exception as e:
    if "unique constraint" in str(e).lower():
        # Another request won the race — fetch their mapping and discard our key
        token_hash = await get_jwt_key_mapping_object(
            jwt_claim_name=virtual_key_claim_field,
            jwt_claim_value=claim_value,
            prisma_client=prisma_client,
        )
        # (optionally delete our orphaned key here)
    else:
        raise

Alternatively, use create_or_update/upsert logic if Prisma supports it for this table.

[greptile-apps]

Inline import violates project style guide

CLAUDE.md states "Avoid imports within methods — place all imports at the top of the file (module-level). The only exception is avoiding circular imports where absolutely necessary." If this import is genuinely required here to break a circular dependency, a brief comment explaining the circular import would make the exception clear to future reviewers.

Suggested change

	from litellm.proxy.management_endpoints.key_management_endpoints import (
	generate_key_helper_fn,
	)
	from litellm.proxy.management_endpoints.key_management_endpoints import (
	generate_key_helper_fn,
	) # local import to avoid circular dependency: user_api_key_auth ← proxy_server → key_management_endpoints

[shivamrawat1]

@greptile re-review now that we have resolved the raised issues

[shivamrawat1]

@greptile re-review with the new commit.

[shivamrawat1]

@review again with new commit