mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(jwt): auto_register_map_existing_key maps JWT to the user's existing virtual key
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng
parent
7c87513dd0
commit
e140e0236a
3 changed files with 304 additions and 35 deletions
|
|
@ -5271,6 +5271,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
"'auto_register': auto-create a virtual key and mapping on first encounter."
|
||||
),
|
||||
)
|
||||
auto_register_map_existing_key: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Only used with unregistered_jwt_client_behavior='auto_register'. When True, the JWT claim is "
|
||||
"mapped to a virtual key the resolved internal user already owns instead of minting a new one. "
|
||||
"If the user owns several, the most recently created non-expired key is chosen. A new key is "
|
||||
"only minted when the user has none."
|
||||
),
|
||||
)
|
||||
routing_overrides: list[JWTRoutingOverride] | None = Field(
|
||||
default=None,
|
||||
description="Optional claim-based routing overrides for JWT-shaped tokens. Matching rules route requests to oauth2 before default JWT flow.",
|
||||
|
|
|
|||
|
|
@ -133,6 +133,7 @@ from litellm.proxy.utils import (
|
|||
normalize_route_for_root_path,
|
||||
)
|
||||
from litellm.repositories.table_repositories import TeamMembershipRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
|
@ -927,6 +928,20 @@ class _PendingAutoRegister(NamedTuple):
|
|||
jwt_issuer: str | None = None
|
||||
|
||||
|
||||
async def _latest_active_key_hash_for_user(prisma_client: PrismaClient, user_id: str) -> str | None:
|
||||
row: Final = await VerificationTokenRepository(prisma_client).table.find_first(
|
||||
where={ # mutable-ok: the prisma where clause contract is a plain dict
|
||||
"user_id": user_id,
|
||||
"OR": [ # mutable-ok: prisma filter literal
|
||||
{"expires": None}, # mutable-ok: prisma filter literal
|
||||
{"expires": {"gt": datetime.now(timezone.utc)}}, # mutable-ok: prisma filter literal
|
||||
],
|
||||
},
|
||||
order={"created_at": "desc"}, # mutable-ok: prisma order literal
|
||||
)
|
||||
return None if row is None else row.token
|
||||
|
||||
|
||||
async def _auto_register_jwt_mapping(
|
||||
virtual_key_claim_field: str,
|
||||
claim_value: str,
|
||||
|
|
@ -945,8 +960,10 @@ async def _auto_register_jwt_mapping(
|
|||
) -> UserAPIKeyAuth | None:
|
||||
"""
|
||||
Auto-register: create a new virtual key + mapping for an unrecognised JWT
|
||||
claim value. ``team_id`` and ``user_id`` must come from a successful
|
||||
``JWTAuthManager.auth_builder`` run — they encode the JWT identity AFTER
|
||||
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``
|
||||
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
|
||||
is stamped with those values so the cached future-request path inherits
|
||||
the same team/user/org limits the auth_builder path would have applied.
|
||||
|
|
@ -962,29 +979,38 @@ async def _auto_register_jwt_mapping(
|
|||
generate_key_helper_fn,
|
||||
)
|
||||
|
||||
# ``table_name="key"`` is required: without it, generate_key_helper_fn
|
||||
# falls into the user-upsert branch (`table_name is None or "user"`) and
|
||||
# attempts to insert into LiteLLM_UserTable with user_id=None, which fails
|
||||
# the NOT NULL @id constraint. Every successful key-creation caller (e.g.
|
||||
# /key/generate) passes table_name="key" explicitly.
|
||||
key_data: Final = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
table_name="key",
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
organization_id=org_id,
|
||||
agent_id=agent_id,
|
||||
metadata={
|
||||
"auto_registered": True,
|
||||
"jwt_claim_field": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
},
|
||||
existing_token_hash: Final = (
|
||||
await _latest_active_key_hash_for_user(prisma_client, user_id)
|
||||
if jwt_handler.litellm_jwtauth.auto_register_map_existing_key and user_id is not None
|
||||
else None
|
||||
)
|
||||
# generate_key_helper_fn returns the plaintext key in "token"; the persisted
|
||||
# row in LiteLLM_VerificationToken uses its hash, so hash here to get the FK
|
||||
# value referenced by LiteLLM_JWTKeyMapping.token.
|
||||
token_hash = hash_token(key_data["token"])
|
||||
minted: Final = existing_token_hash is None
|
||||
if existing_token_hash is not None:
|
||||
token_hash = existing_token_hash
|
||||
else:
|
||||
# ``table_name="key"`` is required: without it, generate_key_helper_fn
|
||||
# falls into the user-upsert branch (`table_name is None or "user"`) and
|
||||
# attempts to insert into LiteLLM_UserTable with user_id=None, which fails
|
||||
# the NOT NULL @id constraint. Every successful key-creation caller (e.g.
|
||||
# /key/generate) passes table_name="key" explicitly.
|
||||
key_data: Final = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
table_name="key",
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
organization_id=org_id,
|
||||
agent_id=agent_id,
|
||||
metadata={ # mutable-ok: GenerateKeyRequest metadata is a plain dict field
|
||||
"auto_registered": True,
|
||||
"jwt_claim_field": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
},
|
||||
)
|
||||
# generate_key_helper_fn returns the plaintext key in "token"; the persisted
|
||||
# row in LiteLLM_VerificationToken uses its hash, so hash here to get the FK
|
||||
# value referenced by LiteLLM_JWTKeyMapping.token.
|
||||
token_hash = hash_token(key_data["token"])
|
||||
|
||||
try:
|
||||
await prisma_client.db.litellm_jwtkeymapping.create(
|
||||
|
|
@ -1011,15 +1037,18 @@ async def _auto_register_jwt_mapping(
|
|||
virtual_key_claim_field,
|
||||
claim_value,
|
||||
)
|
||||
try:
|
||||
await prisma_client.db.litellm_verificationtoken.delete(where={"token": token_hash})
|
||||
except Exception as delete_err:
|
||||
# Don't fail the request if cleanup fails — the orphan is
|
||||
# unmapped and inert. Log so an operator can prune it later.
|
||||
verbose_proxy_logger.warning(
|
||||
"JWT Key Mapping (auto_register): failed to delete orphaned key after race: %s",
|
||||
delete_err,
|
||||
)
|
||||
if minted:
|
||||
try:
|
||||
await prisma_client.db.litellm_verificationtoken.delete(
|
||||
where={"token": token_hash} # mutable-ok: prisma where clause contract is a plain dict
|
||||
)
|
||||
except Exception as delete_err:
|
||||
# Don't fail the request if cleanup fails — the orphan is
|
||||
# unmapped and inert. Log so an operator can prune it later.
|
||||
verbose_proxy_logger.warning(
|
||||
"JWT Key Mapping (auto_register): failed to delete orphaned key after race: %s",
|
||||
delete_err,
|
||||
)
|
||||
token_hash = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
|
|
@ -1049,7 +1078,8 @@ async def _auto_register_jwt_mapping(
|
|||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"JWT Key Mapping (auto_register): created new virtual key for %s='%s'.",
|
||||
"JWT Key Mapping (auto_register): %s virtual key for %s='%s'.",
|
||||
"created new" if minted else "mapped existing",
|
||||
virtual_key_claim_field,
|
||||
claim_value,
|
||||
)
|
||||
|
|
@ -1063,7 +1093,8 @@ async def _auto_register_jwt_mapping(
|
|||
).resolve(hashed_token=token_hash)
|
||||
)
|
||||
if auto_registered_key is not None:
|
||||
auto_registered_key.org_id = org_id
|
||||
if minted:
|
||||
auto_registered_key.org_id = org_id
|
||||
auto_registered_key.end_user_id = end_user_id
|
||||
auto_registered_key.api_key = auto_registered_key.token
|
||||
return auto_registered_key
|
||||
|
|
|
|||
|
|
@ -2090,6 +2090,235 @@ async def test_auto_register_binds_api_key_to_token_hash():
|
|||
assert result.end_user_id == "validated-end-user"
|
||||
|
||||
|
||||
def _auto_register_patches(*, plaintext_key: str | None = "sk-minted-plaintext"):
|
||||
"""The two collaborators every _auto_register_jwt_mapping test patches the same
|
||||
way: key minting (returns {"token": plaintext}) and IdentityStore.resolve."""
|
||||
from litellm.proxy.auth.auth_method import AuthMethod
|
||||
from litellm.proxy.auth.resolvers.models import CredentialRef
|
||||
from litellm.proxy.auth.resolvers.store import IdentityStore
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
resolved_key = UserAPIKeyAuth(
|
||||
token="existing-hash" if plaintext_key is None else hash_token(plaintext_key),
|
||||
user_id="validated-user",
|
||||
team_id="validated-team",
|
||||
org_id="key-own-org",
|
||||
)
|
||||
principal = IdentityStore._principal_from_key(
|
||||
resolved_key,
|
||||
auth_method=AuthMethod.API_KEY,
|
||||
credential_ref=CredentialRef(token_id=resolved_key.token),
|
||||
)
|
||||
return (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"token": plaintext_key},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
|
||||
new_callable=AsyncMock,
|
||||
return_value=principal,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler, **over):
|
||||
kwargs = {
|
||||
"virtual_key_claim_field": "sub",
|
||||
"claim_value": "user1",
|
||||
"jwt_handler": jwt_handler,
|
||||
"prisma_client": prisma_client,
|
||||
"user_api_key_cache": user_api_key_cache,
|
||||
"parent_otel_span": None,
|
||||
"proxy_logging_obj": MagicMock(),
|
||||
"cache_key": "jwt_key_mapping:sub:user1",
|
||||
"team_id": "validated-team",
|
||||
"user_id": "validated-user",
|
||||
"org_id": "jwt-org",
|
||||
"end_user_id": "validated-end-user",
|
||||
}
|
||||
kwargs.update(over)
|
||||
return kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_map_existing_key_reuses_users_latest_key():
|
||||
"""With auto_register_map_existing_key on, the mapping must point at the user's
|
||||
most recently created non-expired key hash and nothing may be minted."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(token="existing-hash")
|
||||
)
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
auto_register_map_existing_key=True,
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
)
|
||||
|
||||
generate_patch, resolve_patch = _auto_register_patches(plaintext_key=None)
|
||||
with generate_patch as generate_key, resolve_patch:
|
||||
result = await _auto_register_jwt_mapping(
|
||||
**_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler)
|
||||
)
|
||||
|
||||
generate_key.assert_not_awaited()
|
||||
|
||||
find_first = prisma_client.db.litellm_verificationtoken.find_first
|
||||
find_first.assert_awaited_once()
|
||||
where = find_first.await_args.kwargs["where"]
|
||||
assert where["user_id"] == "validated-user"
|
||||
assert {"expires": None} in where["OR"], f"non-expired keys must be included: {where}"
|
||||
assert any(isinstance(entry.get("expires"), dict) and "gt" in entry["expires"] for entry in where["OR"]), (
|
||||
f"future-expiring keys must be included: {where}"
|
||||
)
|
||||
assert find_first.await_args.kwargs["order"] == {"created_at": "desc"}
|
||||
|
||||
create_data = prisma_client.db.litellm_jwtkeymapping.create.await_args.kwargs["data"]
|
||||
assert create_data["token"] == "existing-hash"
|
||||
assert create_data["created_by"] == "auto_register"
|
||||
assert user_api_key_cache.async_set_cache.await_args.kwargs["value"] == "existing-hash"
|
||||
assert result is not None
|
||||
assert result.token == "existing-hash"
|
||||
assert result.api_key == "existing-hash"
|
||||
assert result.org_id == "key-own-org"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_map_existing_key_mints_when_user_has_no_key():
|
||||
"""The flag must not leave a keyless user unmapped: with no existing key it
|
||||
falls back to the mint path and maps the claim to the new key's hash."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
auto_register_map_existing_key=True,
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
)
|
||||
|
||||
generate_patch, resolve_patch = _auto_register_patches()
|
||||
with generate_patch as generate_key, resolve_patch:
|
||||
result = await _auto_register_jwt_mapping(
|
||||
**_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler)
|
||||
)
|
||||
|
||||
generate_key.assert_awaited_once()
|
||||
create_data = prisma_client.db.litellm_jwtkeymapping.create.await_args.kwargs["data"]
|
||||
assert create_data["token"] == hash_token("sk-minted-plaintext")
|
||||
assert result is not None
|
||||
assert result.token == hash_token("sk-minted-plaintext")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_default_never_looks_up_existing_keys():
|
||||
"""Without the flag the behavior is unchanged: no verification-token lookup at
|
||||
all, a fresh key is always minted."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(token="existing-hash")
|
||||
)
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
|
||||
|
||||
generate_patch, resolve_patch = _auto_register_patches()
|
||||
with generate_patch as generate_key, resolve_patch:
|
||||
await _auto_register_jwt_mapping(**_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler))
|
||||
|
||||
prisma_client.db.litellm_verificationtoken.find_first.assert_not_awaited()
|
||||
generate_key.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_map_existing_key_race_loser_keeps_reused_key():
|
||||
"""A reused key is not ours to delete: when the unique-constraint race is lost,
|
||||
the pre-existing user key must survive and the winner's mapping wins."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(token="existing-hash")
|
||||
)
|
||||
prisma_client.db.litellm_verificationtoken.delete = AsyncMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(side_effect=Exception("Unique constraint failed (P2002)"))
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
auto_register_map_existing_key=True,
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
)
|
||||
|
||||
generate_patch, resolve_patch = _auto_register_patches(plaintext_key=None)
|
||||
with (
|
||||
generate_patch,
|
||||
resolve_patch,
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_jwt_key_mapping_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value="winner-hash",
|
||||
),
|
||||
):
|
||||
await _auto_register_jwt_mapping(**_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler))
|
||||
|
||||
prisma_client.db.litellm_verificationtoken.delete.assert_not_awaited()
|
||||
assert user_api_key_cache.async_set_cache.await_args.kwargs["value"] == "winner-hash"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_map_existing_key_user_id_none_mints():
|
||||
"""With no resolved user there is no key to reuse; the flag must not skip
|
||||
minting, and the lookup must not run."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(token="existing-hash")
|
||||
)
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
auto_register_map_existing_key=True,
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
)
|
||||
|
||||
generate_patch, resolve_patch = _auto_register_patches()
|
||||
with generate_patch as generate_key, resolve_patch:
|
||||
await _auto_register_jwt_mapping(
|
||||
**_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler, user_id=None)
|
||||
)
|
||||
|
||||
prisma_client.db.litellm_verificationtoken.find_first.assert_not_awaited()
|
||||
generate_key.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("active", [True, False])
|
||||
async def test_auto_register_first_request_propagates_user_email(active: bool) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue