feat(auth): add link-only JWT mapping and remove route eligibility filtering (#49956)

* feat(auth): allow skipping JWT key route eligibility check

* feat(auth): support link-only JWT key mapping

* fix(auth): model JWT mapping rejection as a failure value

* feat(auth): select existing JWT keys without route filtering
This commit is contained in:
mrinal-berri 2026-11-03 16:24:51 -08:00 • committed by GitHub
parent ce73cf8b9d
commit 7fa438df06
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 406 additions and 44 deletions

View file

@ -4796,11 +4796,14 @@ class UnregisteredJWTClientBehavior(str, enum.Enum):
- 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.
- link_only: Map an eligible existing user key or return HTTP 403, without
creating a key. Proxy admins authenticate without a mapped key.
"""
FALLBACK_TEAM_MAPPING = "fallback_team_mapping"
REJECT = "reject"
AUTO_REGISTER = "auto_register"
LINK_ONLY = "link_only"
class JWTIssuerConfig(LiteLLMBaseModel):
@ -5038,7 +5041,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
"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."
"'auto_register': auto-create a virtual key and mapping on first encounter. "
"'link_only': map an eligible existing user key or return HTTP 403, without creating a key. "
"link_only ignores auto_register_map_existing_key; user provisioning still follows user_id_upsert. "
"Existing-key selection ignores allowed_routes; JWTs inherit the selected key's permissions, "
"including management routes, with request-time key restrictions still enforced. "
"Unmapped proxy admins authenticate directly in auto_register and link_only modes."
),
)
auto_register_map_existing_key: bool = Field(
@ -5048,8 +5056,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
"field is the user_id_jwt_field or user_email_jwt_field, the JWT claim is mapped to a virtual key the "
"JWT-resolved user already owns instead of minting a new one. If the user owns several, the most recently created key in the "
"JWT-resolved team (or with no team when the JWT resolves none) is chosen among keys that never "
"expire, are not blocked, are not Admin UI session keys, were not minted by auto_register, and "
"have no allowed_routes or include llm_api_routes. Otherwise a new key is minted as usual."
"expire, are not blocked, are not Admin UI session keys, and were not minted by auto_register. "
"allowed_routes does not affect key selection. JWTs inherit the selected key's permissions, "
"including management routes; request-time key restrictions still apply. "
"Otherwise a new key is minted as usual."
),
)
routing_overrides: list[JWTRoutingOverride] | None = Field(

View file

@ -13,6 +13,7 @@ import re
import secrets
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from enum import Enum
from types import MappingProxyType
from typing import Final, NamedTuple, Protocol, Union, cast
@ -1576,10 +1577,14 @@ _JWT_AUTH_DISABLED_HINT = (
)
class _JWTKeyMappingFailure(Enum):
NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied."
class _PendingAutoRegister(NamedTuple):
"""
Signal returned by ``_resolve_jwt_to_virtual_key`` when the JWT's claim is
unmapped and ``unregistered_jwt_client_behavior`` is AUTO_REGISTER.
unmapped and ``unregistered_jwt_client_behavior`` is AUTO_REGISTER or LINK_ONLY.
The caller MUST run standard ``JWTAuthManager.auth_builder`` to apply RBAC,
scope mappings, ``custom_validate``, and ``user_allowed_email_domain``
@ -1593,23 +1598,22 @@ class _PendingAutoRegister(NamedTuple):
claim_value: str
cache_key: str
jwt_issuer: str | None = None
behavior: UnregisteredJWTClientBehavior = UnregisteredJWTClientBehavior.AUTO_REGISTER
def _claim_identifies_user(jwt_handler: JWTHandler, claim_field: str, jwt_issuer: str | None) -> bool:
if not jwt_handler.litellm_jwtauth.auto_register_map_existing_key:
return False
if jwt_handler.litellm_jwtauth.is_user_identity_claim(claim_field, jwt_issuer):
return True
verbose_proxy_logger.warning(
"JWT Key Mapping (auto_register_map_existing_key): claim '%s' is not the user_id or user_email JWT field "
"and may be shared by several users, so a new key is minted instead of reusing one the user owns.",
"JWT Key Mapping: claim '%s' is not the user_id or user_email JWT field "
"and may be shared by several users, so existing user keys cannot be linked automatically.",
claim_field,
)
return False
async def _reusable_key_hash_for_user(prisma_client: PrismaClient, user_id: str, team_id: str | None) -> str | None:
key: Final = await VerificationTokenRepository(prisma_client).find_newest_reusable_llm_api_key(user_id, team_id)
key: Final = await VerificationTokenRepository(prisma_client).find_newest_reusable_key(user_id, team_id)
return None if key is None else key.token
@ -1628,11 +1632,13 @@ async def _auto_register_jwt_mapping(
org_id: str | None = None,
end_user_id: str | None = None,
agent_id: str | None = None,
) -> UserAPIKeyAuth | None:
behavior: UnregisteredJWTClientBehavior = UnregisteredJWTClientBehavior.AUTO_REGISTER,
) -> UserAPIKeyAuth | _JWTKeyMappingFailure | None:
"""
Auto-register: create a new virtual key + mapping for an unrecognised JWT
claim value, or point the mapping at a key the resolved user already owns
when ``auto_register_map_existing_key`` is set. ``team_id`` and ``user_id``
when ``auto_register_map_existing_key`` is set. LINK_ONLY always searches
for an existing key and rejects when none is eligible. ``team_id`` and ``user_id``
must come from a successful ``JWTAuthManager.auth_builder`` run — they
encode the JWT identity AFTER
RBAC/scope/custom_validate/email-domain policy has been enforced. The key
@ -1650,11 +1656,21 @@ async def _auto_register_jwt_mapping(
generate_key_helper_fn,
)
link_only: Final = behavior == UnregisteredJWTClientBehavior.LINK_ONLY
reuse_existing: Final = link_only or jwt_handler.litellm_jwtauth.auto_register_map_existing_key
existing_token_hash: Final = (
await _reusable_key_hash_for_user(prisma_client, user_id, team_id)
if user_id is not None and _claim_identifies_user(jwt_handler, virtual_key_claim_field, jwt_issuer)
await _reusable_key_hash_for_user(
prisma_client,
user_id,
team_id,
)
if reuse_existing
and user_id is not None
and _claim_identifies_user(jwt_handler, virtual_key_claim_field, jwt_issuer)
else None
)
if link_only and existing_token_hash is None:
return _JWTKeyMappingFailure.NO_ELIGIBLE_KEY
minted: Final = existing_token_hash is None
if existing_token_hash is not None:
token_hash = existing_token_hash
@ -1691,8 +1707,8 @@ async def _auto_register_jwt_mapping(
"jwt_claim_name": virtual_key_claim_field,
"jwt_claim_value": claim_value,
"token": token_hash,
"created_by": "auto_register",
"updated_by": "auto_register",
"created_by": behavior.value,
"updated_by": behavior.value,
}
)
except Exception as e:
@ -1704,8 +1720,8 @@ async def _auto_register_jwt_mapping(
# the DB and the cleartext is in memory on this request. Delete it
# so orphans don't accumulate under sustained concurrency.
verbose_proxy_logger.debug(
"JWT Key Mapping (auto_register): unique conflict on create — "
"deleting orphaned virtual key and fetching winner's mapping for %s='%s'.",
"JWT Key Mapping (%s): unique conflict on create, fetching winner's mapping for %s='%s'.",
behavior.value,
virtual_key_claim_field,
claim_value,
)
@ -1735,7 +1751,7 @@ async def _auto_register_jwt_mapping(
raise HTTPException(
status_code=503,
detail=(
"JWT Key Mapping: AUTO_REGISTER race resolution failed — "
f"JWT Key Mapping: {behavior.value.upper()} race resolution failed — "
"winner's mapping was concurrently removed. Retry the request."
),
)
@ -1749,7 +1765,8 @@ async def _auto_register_jwt_mapping(
)
verbose_proxy_logger.info(
"JWT Key Mapping (auto_register): %s virtual key for %s='%s'.",
"JWT Key Mapping (%s): %s virtual key for %s='%s'.",
behavior.value,
"created new" if minted else "mapped existing",
virtual_key_claim_field,
claim_value,
@ -1821,7 +1838,7 @@ async def _resolve_jwt_to_virtual_key(
- ``UserAPIKeyAuth``: a resolved virtual key (cache hit or DB hit). The
caller may use this directly; JWT policy has been enforced previously
(at key-creation time or, for cached results, before caching).
- ``_PendingAutoRegister``: claim is unmapped and behavior is AUTO_REGISTER.
- ``_PendingAutoRegister``: claim is unmapped and behavior is AUTO_REGISTER or LINK_ONLY.
The caller MUST run ``JWTAuthManager.auth_builder`` to enforce JWT
policy (RBAC, scope, custom_validate, email-domain), then invoke
``_auto_register_jwt_mapping`` with the validated team_id/user_id.
@ -1829,7 +1846,7 @@ async def _resolve_jwt_to_virtual_key(
The caller falls through to standard team-based JWT auth (which itself
enforces full JWT policy via auth_builder).
- Raises HTTPException: REJECT policy hit, missing claim under
REJECT/AUTO_REGISTER, or other policy violations.
REJECT/AUTO_REGISTER/LINK_ONLY, or other policy violations.
"""
raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM)
normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None
@ -1856,6 +1873,7 @@ async def _resolve_jwt_to_virtual_key(
if behavior in (
UnregisteredJWTClientBehavior.REJECT,
UnregisteredJWTClientBehavior.AUTO_REGISTER,
UnregisteredJWTClientBehavior.LINK_ONLY,
):
raise HTTPException(
status_code=403,
@ -1887,7 +1905,7 @@ async def _resolve_jwt_to_virtual_key(
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 behavior in (UnregisteredJWTClientBehavior.AUTO_REGISTER, UnregisteredJWTClientBehavior.LINK_ONLY):
# Stale sentinel written under a prior fallback_team_mapping config —
# evict it and defer auto-register to after auth_builder runs. Raise
# the same 500 as the fresh-path AUTO_REGISTER branch when there is
@ -1897,7 +1915,7 @@ async def _resolve_jwt_to_virtual_key(
raise HTTPException(
status_code=500,
detail=(
"JWT Key Mapping: AUTO_REGISTER requires a database connection. "
f"JWT Key Mapping: {behavior.value.upper()} requires a database connection. "
"Configure a database or change unregistered_jwt_client_behavior."
),
)
@ -1907,6 +1925,7 @@ async def _resolve_jwt_to_virtual_key(
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
behavior=behavior,
)
return None
elif cached_mapping is not None:
@ -1967,12 +1986,12 @@ async def _resolve_jwt_to_virtual_key(
detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.",
)
if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER:
if behavior in (UnregisteredJWTClientBehavior.AUTO_REGISTER, UnregisteredJWTClientBehavior.LINK_ONLY):
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=(
"JWT Key Mapping: AUTO_REGISTER requires a database connection. "
f"JWT Key Mapping: {behavior.value.upper()} requires a database connection. "
"Configure a database or change unregistered_jwt_client_behavior."
),
)
@ -1985,6 +2004,7 @@ async def _resolve_jwt_to_virtual_key(
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
behavior=behavior,
)
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
@ -2477,10 +2497,13 @@ async def user_api_key_auth_builder(
org_id=org_id,
end_user_id=end_user_id,
agent_id=agent_id,
behavior=pending_auto_register.behavior,
)
if pending_auto_register is not None and prisma_client is not None
else None
)
if isinstance(auto_registered, _JWTKeyMappingFailure):
raise HTTPException(status_code=403, detail=auto_registered.value)
if auto_registered is not None:
auto_registered.jwt_claims = jwt_claims
auto_registered.user_email = user_email
@ -2496,7 +2519,13 @@ async def user_api_key_auth_builder(
falls_through_to_key_checks: Final = (
auto_registered is not None
and jwt_handler.litellm_jwtauth.auto_register_map_existing_key
and (
jwt_handler.litellm_jwtauth.auto_register_map_existing_key
or (
pending_auto_register is not None
and pending_auto_register.behavior == UnregisteredJWTClientBehavior.LINK_ONLY
)
)
and master_key is not None
)
if not falls_through_to_key_checks:

View file

@ -124,9 +124,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"user_id": user_id})
return self._to_model_list(records)
async def find_newest_reusable_llm_api_key(
self, user_id: str, team_id: str | None
) -> LiteLLM_VerificationToken | None:
async def find_newest_reusable_key(self, user_id: str, team_id: str | None) -> LiteLLM_VerificationToken | None:
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(
where={
"user_id": user_id,
@ -140,12 +138,6 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
{"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
]
},
{
"OR": [
{"allowed_routes": {"is_empty": True}},
{"allowed_routes": {"has": "llm_api_routes"}},
]
},
],
},
order={"created_at": "desc"},

View file

@ -7,13 +7,15 @@ from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from hashlib import sha256
from pathlib import Path
from typing import Final
from typing import Final, Literal
import httpx
import jwt
import pytest
import yaml
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import JsonValue
from redis import Redis
from tests.integration._support.client import Gateway, eventually, string_value
from tests.integration._support.database import read_rows
@ -31,7 +33,15 @@ def _hash(key: str) -> str:
return sha256(key.encode()).hexdigest()
def _config(directory: Path, claim_field: str) -> Path:
def _config(
directory: Path,
claim_field: str,
*,
behavior: Literal["auto_register", "link_only", "reject"] = "auto_register",
map_existing_key: bool = True,
user_id_upsert: bool = False,
issuers: tuple[dict[str, JsonValue], ...] = (),
) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"] = {
**config["general_settings"],
@ -40,8 +50,10 @@ def _config(directory: Path, claim_field: str) -> Path:
"user_id_jwt_field": "sub",
"user_email_jwt_field": "email",
"virtual_key_claim_field": claim_field,
"unregistered_jwt_client_behavior": "auto_register",
"auto_register_map_existing_key": True,
"unregistered_jwt_client_behavior": behavior,
"auto_register_map_existing_key": map_existing_key,
"user_id_upsert": user_id_upsert,
"issuers": list(issuers),
},
}
path: Final = directory / f"jwt_map_existing_key_{claim_field}.yaml"
@ -107,7 +119,7 @@ def _billed_key(response: httpx.Response) -> str:
return string_value(rows[0]["api_key"])
def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_ineligible_newer_key(
def test_first_jwt_call_reuses_the_newest_durable_key_and_skips_every_ineligible_newer_key(
gateway: Gateway, tmp_path: Path
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
@ -118,7 +130,6 @@ def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_inelig
skipped: Final = {
"older_durable": older_durable,
"expiring": scenario.key(user_id=user, duration="1h"),
"management_only": scenario.key(user_id=user, allowed_routes=["management_routes"]),
"auto_registered_look_alike": scenario.key(user_id=user, metadata={"auto_registered": True}),
"other_team": scenario.key(user_id=user, team_id=scenario.team()),
"blocked": scenario.key(user_id=user),
@ -140,6 +151,84 @@ def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_inelig
assert _billed_key(response) == _hash(durable)
@pytest.mark.parametrize(
("behavior", "map_existing_key"), [("auto_register", True), ("link_only", False), ("link_only", True)]
)
def test_route_restricted_keys_are_linked_without_bypassing_other_key_checks(
gateway: Gateway,
tmp_path: Path,
behavior: Literal["auto_register", "link_only"],
map_existing_key: bool,
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
scenario.key(user_id=user)
restricted: Final = scenario.key(user_id=user, allowed_routes=["/v1/chat/completions"])
scenario.key(user_id=user, duration="1h")
scenario.key(user_id=user, metadata={"auto_registered": True})
scenario.key(user_id=user, team_id=scenario.team())
blocked: Final = scenario.key(user_id=user)
gateway.post("/key/block", {"key": blocked})
keys_before: Final = _user_key_hashes(user)
token: Final = _token(private_key, user)
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(
tmp_path,
"sub",
behavior=behavior,
map_existing_key=map_existing_key,
),
) as candidate:
response: Final = _chat(candidate, model, token)
repeated: Final = _chat(candidate, model, token)
restricted_route: Final = candidate.request("GET", "/v1/models", key=token)
assert response.status_code == 200, response.text
assert repeated.status_code == 200, repeated.text
assert restricted_route.status_code == 403, restricted_route.text
assert _mapped_token("sub", user) == _hash(restricted)
assert _user_key_hashes(user) == keys_before, "a new key was minted instead of linking an existing key"
assert _billed_key(response) == _hash(restricted)
assert _billed_key(repeated) == _hash(restricted)
@pytest.mark.parametrize("behavior", ["auto_register", "link_only"])
def test_management_route_keys_are_linked_by_default_and_keep_their_route_permissions(
gateway: Gateway, tmp_path: Path, behavior: Literal["auto_register", "link_only"]
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
target: Final = f"integration-{uuid.uuid4().hex}"
scenario.key(user_id=user)
management_key: Final = scenario.key(user_id=user, allowed_routes=["/user/new"])
keys_before: Final = _user_key_hashes(user)
token: Final = _token(private_key, user)
with owned_proxy(
gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks_url}, config=_config(tmp_path, "sub", behavior=behavior)
) as candidate:
denied: Final = _chat(candidate, model, token)
allowed: Final = candidate.request(
"POST",
"/user/new",
{"user_id": target, "auto_create_key": False, "user_role": "internal_user"},
key=token,
)
assert allowed.status_code == 200, allowed.text
scenario.cleanups.callback(scenario.delete_user, target)
repeated: Final = _chat(candidate, model, token)
assert denied.status_code == 403, denied.text
assert repeated.status_code == 403, repeated.text
assert allowed.json()["user_id"] == target
assert _mapped_token("sub", user) == _hash(management_key)
assert _user_key_hashes(user) == keys_before
def test_user_matched_by_email_instead_of_sub_still_reuses_their_existing_key(gateway: Gateway, tmp_path: Path) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
@ -183,8 +272,9 @@ def test_shared_client_claim_never_maps_a_second_user_onto_the_first_users_perso
assert read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (_hash(personal),)) == []
def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviving_mapped_key(
gateway: Gateway, tmp_path: Path
@pytest.mark.parametrize("behavior", ["auto_register", "link_only"])
def test_concurrent_first_jwt_calls_both_succeed_on_one_surviving_mapped_key(
gateway: Gateway, tmp_path: Path, behavior: Literal["auto_register", "link_only"]
) -> None:
writer_url: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL") or os.environ["DATABASE_URL"]
with (
@ -194,6 +284,7 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv
):
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
existing: Final = scenario.key(user_id=user) if behavior == "link_only" else None
token: Final = _token(private_key, user)
overrides: Final = {
"JWT_PUBLIC_KEY_URL": jwks_url,
@ -202,7 +293,12 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv
}
with (
owned_proxy(gateway, tmp_path, overrides, config=_config(tmp_path, "sub")) as candidate,
owned_proxy(
gateway,
tmp_path,
overrides,
config=_config(tmp_path, "sub", behavior=behavior, map_existing_key=False),
) as candidate,
ThreadPoolExecutor(max_workers=1) as pool,
):
held_call: Final = pool.submit(_chat, candidate, model, token)
@ -215,5 +311,240 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv
assert held.status_code == 200, held.text
keys: Final = _user_key_hashes(user)
assert len(keys) == 1, keys
if existing is not None:
assert keys == frozenset({_hash(existing)})
assert _mapped_token("sub", user) in keys
assert (_billed_key(held), _billed_key(racing)) == (_mapped_token("sub", user),) * 2
@pytest.mark.parametrize("ineligible_keys", [False, True])
def test_link_only_rejects_without_writes_then_links_a_newly_available_key(
gateway: Gateway, tmp_path: Path, ineligible_keys: bool
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
if ineligible_keys:
scenario.key(user_id=user, duration="1h")
scenario.key(user_id=user, metadata={"auto_registered": True})
scenario.key(user_id=user, team_id=scenario.team())
blocked: Final = scenario.key(user_id=user)
gateway.post("/key/block", {"key": blocked})
keys_before: Final = _user_key_hashes(user)
token: Final = _token(private_key, user)
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
) as candidate:
first: Final = _chat(candidate, model, token)
repeated: Final = _chat(candidate, model, token)
assert first.status_code == 403, first.text
assert repeated.status_code == 403, repeated.text
assert _user_key_hashes(user) == keys_before
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == []
available: Final = scenario.key(user_id=user)
accepted: Final = _chat(candidate, model, token)
assert accepted.status_code == 200, accepted.text
assert _mapped_token("sub", user) == _hash(available)
assert _user_key_hashes(user) == keys_before | {_hash(available)}
assert _billed_key(accepted) == _hash(available)
@pytest.mark.parametrize("claim_field", ["client_id", "email"])
def test_link_only_rejects_shared_or_missing_claims_without_linking_a_personal_key(
gateway: Gateway, tmp_path: Path, claim_field: str
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
key: Final = scenario.key(user_id=user)
shared_claim: Final = f"integration-client-{uuid.uuid4().hex}"
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, claim_field, behavior="link_only", map_existing_key=False),
) as candidate:
response: Final = _chat(candidate, model, _token(private_key, user, client_id=shared_claim))
assert response.status_code == 403, response.text
assert _user_key_hashes(user) == frozenset({_hash(key)})
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE token=%s', (_hash(key),)) == []
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (shared_claim,)) == []
@pytest.mark.parametrize("restriction", ["models", "allowed_routes"])
def test_link_only_enforces_key_permissions_on_the_first_and_cached_requests(
gateway: Gateway, tmp_path: Path, restriction: str
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
allowed: Final = scenario.model()
denied: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
key: Final = scenario.key(
user_id=user,
**({"models": [allowed]} if restriction == "models" else {"allowed_routes": ["/v1/models"]}),
)
token: Final = _token(private_key, user)
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
) as candidate:
first: Final = _chat(candidate, denied, token)
repeated: Final = _chat(candidate, denied, token)
control: Final = (
_chat(candidate, allowed, token)
if restriction == "models"
else candidate.request("GET", "/v1/models", key=token)
)
assert first.status_code == 403, first.text
assert repeated.status_code == 403, repeated.text
assert control.status_code == 200, control.text
assert _mapped_token("sub", user) == _hash(key)
assert _user_key_hashes(user) == frozenset({_hash(key)})
def test_link_only_does_not_replace_an_existing_mapping_to_a_blocked_key(gateway: Gateway, tmp_path: Path) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
blocked: Final = scenario.key(user_id=user)
available: Final = scenario.key(user_id=user)
gateway.post("/jwt/key/mapping/new", {"jwt_claim_name": "sub", "jwt_claim_value": user, "key": blocked})
gateway.post("/key/block", {"key": blocked})
token: Final = _token(private_key, user)
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
) as candidate:
first: Final = _chat(candidate, model, token)
repeated: Final = _chat(candidate, model, token)
assert first.status_code == 401, first.text
assert repeated.status_code == 401, repeated.text
assert _mapped_token("sub", user) == _hash(blocked)
assert _user_key_hashes(user) == frozenset({_hash(blocked), _hash(available)})
def test_link_only_exempts_proxy_admins_but_rechecks_the_role_after_a_cached_exemption(
gateway: Gateway, tmp_path: Path
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
admin: Final = _token(private_key, user, scope="litellm_proxy_admin")
ordinary: Final = _token(private_key, user)
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
) as candidate:
first: Final = candidate.request("GET", "/v1/models", key=admin)
repeated: Final = candidate.request("GET", "/v1/models", key=admin)
denied: Final = candidate.request("GET", "/v1/models", key=ordinary)
forbidden_route: Final = _chat(candidate, model, admin)
assert first.status_code == 200, first.text
assert repeated.status_code == 200, repeated.text
assert denied.status_code == 403, denied.text
assert forbidden_route.status_code == 401, forbidden_route.text
assert _user_key_hashes(user) == frozenset()
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == []
@pytest.mark.parametrize("user_id_upsert", [False, True])
def test_link_only_keeps_user_provisioning_independent_of_key_creation(
gateway: Gateway, tmp_path: Path, user_id_upsert: bool
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
subject: Final = f"integration-link-only-{uuid.uuid4().hex}"
try:
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(
tmp_path, "sub", behavior="link_only", map_existing_key=False, user_id_upsert=user_id_upsert
),
) as candidate:
response: Final = _chat(candidate, model, _token(private_key, subject))
assert response.status_code == (403 if user_id_upsert else 401), response.text
users: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (subject,))
assert users == ([{"user_id": subject}] if user_id_upsert else [])
assert _user_key_hashes(subject) == frozenset()
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (subject,)) == []
finally:
if read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (subject,)):
scenario.delete_user(subject)
@pytest.mark.parametrize("issuer_behavior", ["auto_register", "link_only"])
def test_issuer_specific_registration_mode_overrides_the_global_mode(
gateway: Gateway, tmp_path: Path, issuer_behavior: Literal["auto_register", "link_only"]
) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
issuer: Final = f"https://integration.example.com/{uuid.uuid4().hex}"
token: Final = _token(private_key, user, iss=issuer, aud="integration")
path: Final = _config(
tmp_path,
"sub",
behavior="auto_register" if issuer_behavior == "link_only" else "link_only",
map_existing_key=False,
issuers=(
{
"issuer": issuer,
"jwks_url": jwks_url,
"audience": "integration",
"user_id_jwt_field": "sub",
"virtual_key_claim_field": "sub",
"unregistered_jwt_client_behavior": issuer_behavior,
},
),
)
with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks_url}, config=path) as candidate:
response: Final = _chat(candidate, model, token)
assert response.status_code == (403 if issuer_behavior == "link_only" else 200), response.text
if issuer_behavior == "link_only":
assert _user_key_hashes(user) == frozenset()
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == []
else:
assert len(_user_key_hashes(user)) == 1
assert _mapped_token("sub", user) in _user_key_hashes(user)
assert _billed_key(response) == _mapped_token("sub", user)
def test_link_only_links_after_a_cached_rejection_from_the_previous_mode(gateway: Gateway, tmp_path: Path) -> None:
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user(user_role="internal_user")
key: Final = scenario.key(user_id=user)
token: Final = _token(private_key, user)
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, "sub", behavior="reject", map_existing_key=False),
) as rejecting:
denied: Final = _chat(rejecting, model, token)
assert denied.status_code == 403, denied.text
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
cached: Final = cache.get(f"jwt_key_mapping:sub:{user}")
assert cached is not None and b"__NO_MAPPING__" in cached, cached
with owned_proxy(
gateway,
tmp_path,
{"JWT_PUBLIC_KEY_URL": jwks_url},
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
) as linking:
accepted: Final = _chat(linking, model, token)
assert accepted.status_code == 200, accepted.text
assert _mapped_token("sub", user) == _hash(key)
assert _user_key_hashes(user) == frozenset({_hash(key)})
assert _billed_key(accepted) == _hash(key)