scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

fix(jwt): scope auto_register_map_existing_key reuse to the JWT-resolved team

Browse source 
Only reuse a key whose team_id matches the team auth_builder resolved for
the JWT (no team matches no team), so a personal key can no longer bypass
the resolved team's model and budget limits.

With the flag on, the first JWT request now falls through to the same
virtual-key checks later mapped requests get, instead of returning early,
so a reused key's own limits apply from request one rather than 200 then
403. Flag off keeps the early return unchanged.
This commit is contained in:
Yuneng Jiang 2026-09-23 18:23:31 -07:00
parent 449c176ef8
commit 478de14bb5
No known key found for this signature in database
4 changed files with 211 additions and 70 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/_types.py
View file

@ -5303,9 +5303,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
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 key that can call LLM routes is chosen: "
"not blocked, not expired, not an Admin UI session key, and with no route restriction other than "
"llm_api_routes. A new key is only minted when the user has no such key."
"If the user owns several, the most recently created key in the JWT-resolved team (or with no "
"team when the JWT resolves none) that can call LLM routes is chosen: not blocked, not expired, "
"not an Admin UI session key, and with no route restriction other than llm_api_routes. A new key "
"is only minted when the user has no such key."
),
)
routing_overrides: list[JWTRoutingOverride] | None = Field(

138
litellm/proxy/auth/user_api_key_auth.py
View file

@ -930,8 +930,10 @@ 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).find_latest_llm_api_row_by_user_id(user_id)
async def _latest_active_key_hash_for_user(
prisma_client: PrismaClient, user_id: str, team_id: str | None
) -> str | None:
row: Final = await VerificationTokenRepository(prisma_client).find_latest_llm_api_row_by_user_id(user_id, team_id)
return None if row is None else row.token
@ -973,7 +975,7 @@ async def _auto_register_jwt_mapping(
)
existing_token_hash: Final = (
await _latest_active_key_hash_for_user(prisma_client, user_id)
await _latest_active_key_hash_for_user(prisma_client, user_id, team_id)
if jwt_handler.litellm_jwtauth.auto_register_map_existing_key and user_id is not None
else None
)
@ -1772,8 +1774,8 @@ async def _user_api_key_auth_builder(
# mapping + virtual key from the *validated* identity, then
# replace valid_token with the new key so downstream checks
# use the key-scoped path.
if pending_auto_register is not None and prisma_client is not None:
auto_registered: Final = await _auto_register_jwt_mapping(
auto_registered: Final = (
await _auto_register_jwt_mapping(
virtual_key_claim_field=pending_auto_register.claim_field,
claim_value=pending_auto_register.claim_value,
jwt_handler=jwt_handler,
@ -1789,72 +1791,76 @@ async def _user_api_key_auth_builder(
end_user_id=end_user_id,
agent_id=agent_id,
)
if auto_registered is not None:
auto_registered.jwt_claims = jwt_claims
auto_registered.user_email = user_email
# The auto-registered token is built from the new key's
# columns, which carry no user budget. Carry over the
# already-loaded user row rather than re-reading it, or
# the budget check below has nothing to enforce.
auto_registered.user_model_max_budget = (
user_object.model_max_budget if user_object is not None else None
)
valid_token = auto_registered
api_key = valid_token.token or ""
# Check if model has zero cost - if so, skip all budget checks
model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
if pending_auto_register is not None and prisma_client is not None
else None
)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
if skip_budget_checks:
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
# Fetch project object for JWT path if project_id is set
_jwt_project_obj = None
if valid_token.project_id is not None:
_jwt_project_obj = await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
if auto_registered is not None:
auto_registered.jwt_claims = jwt_claims
auto_registered.user_email = user_email
# The auto-registered token is built from the new key's
# columns, which carry no user budget. Carry over the
# already-loaded user row rather than re-reading it, or
# the budget check below has nothing to enforce.
auto_registered.user_model_max_budget = (
user_object.model_max_budget if user_object is not None else None
)
if _jwt_project_obj is not None:
valid_token.project_metadata = _jwt_project_obj.metadata
valid_token.project_alias = _jwt_project_obj.project_alias
valid_token = auto_registered
api_key = valid_token.token or ""
# JWT auth returns here rather than falling through to the
# virtual-key checks below, so the user's per-model budget
# has to be enforced on this path too. Without it the
# post-call increment still charges the counter and nothing
# ever reads it, which is worse than not tracking at all.
# Guarded by the same flag the virtual-key path uses, or a
# zero-cost model would be refused here and allowed there,
# while the log above claims all budget checks were skipped.
if not skip_budget_checks:
await _check_user_model_budget(
valid_token=cast(UserAPIKeyAuth, valid_token),
model_max_budget_limiter=model_max_budget_limiter,
models=_get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
),
if auto_registered is None or not jwt_handler.litellm_jwtauth.auto_register_map_existing_key:
# Check if model has zero cost - if so, skip all budget checks
model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
return cast(UserAPIKeyAuth, valid_token)
skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
if skip_budget_checks:
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
# Fetch project object for JWT path if project_id is set
_jwt_project_obj = None
if valid_token.project_id is not None:
_jwt_project_obj = await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if _jwt_project_obj is not None:
valid_token.project_metadata = _jwt_project_obj.metadata
valid_token.project_alias = _jwt_project_obj.project_alias
# JWT auth returns here rather than falling through to the
# virtual-key checks below, so the user's per-model budget
# has to be enforced on this path too. Without it the
# post-call increment still charges the counter and nothing
# ever reads it, which is worse than not tracking at all.
# Guarded by the same flag the virtual-key path uses, or a
# zero-cost model would be refused here and allowed there,
# while the log above claims all budget checks were skipped.
if not skip_budget_checks:
await _check_user_model_budget(
valid_token=cast(UserAPIKeyAuth, valid_token),
model_max_budget_limiter=model_max_budget_limiter,
models=_get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
),
)
return cast(UserAPIKeyAuth, valid_token)
#### ELSE ####
## CHECK PASS-THROUGH ENDPOINTS ##

5
litellm/repositories/verification_token_repository.py
View file

@ -124,10 +124,13 @@ 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_latest_llm_api_row_by_user_id(self, user_id: str) -> "PrismaVerificationToken | None":
async def find_latest_llm_api_row_by_user_id(
self, user_id: str, team_id: str | None
) -> "PrismaVerificationToken | None":
row: Final = await self.table.find_first(
where={ # mutable-ok: the prisma where clause contract is a plain dict
"user_id": user_id,
"team_id": team_id,
"AND": [ # mutable-ok: prisma filter literal
{"OR": [{"blocked": False}, {"blocked": None}]}, # mutable-ok: prisma filter literal
{ # mutable-ok: prisma filter literal

131
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -2236,6 +2236,39 @@ async def test_auto_register_map_existing_key_skips_keys_that_cannot_call_llm_ro
], f"only unrestricted or llm_api_routes keys may be reused: {where}"
@pytest.mark.asyncio
@pytest.mark.parametrize("resolved_team_id", ["validated-team", None])
async def test_auto_register_map_existing_key_only_reuses_keys_in_the_jwt_resolved_team(
resolved_team_id: str | None,
) -> None:
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, resolve_patch:
await _auto_register_jwt_mapping(
**_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler, team_id=resolved_team_id)
)
where = prisma_client.db.litellm_verificationtoken.find_first.await_args.kwargs["where"]
assert "team_id" in where, f"reuse must be scoped to the JWT-resolved team: {where}"
assert where["team_id"] == resolved_team_id
@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
@ -2364,6 +2397,104 @@ async def test_auto_register_map_existing_key_user_id_none_mints():
generate_key.assert_awaited_once()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("reused_key_models", "expect_denied"),
[(["some-other-model"], True), ([], False)],
)
async def test_auto_register_map_existing_key_first_request_runs_key_checks(
reused_key_models: list[str], expect_denied: bool
) -> None:
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
user_api_key_cache = DualCache()
prisma_client = MagicMock()
jwt_handler = MagicMock()
jwt_handler.is_jwt.return_value = True
jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1"})
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub",
virtual_key_mapping_cache_ttl=300,
auto_register_map_existing_key=True,
)
reused_key = UserAPIKeyAuth(
token="hashed-existing-key",
api_key="hashed-existing-key",
user_id="validated-user",
models=reused_key_models,
)
mock_jwt_result = {
"is_proxy_admin": False,
"team_object": None,
"user_object": LiteLLM_UserTable(user_id="validated-user", user_role="internal_user"),
"end_user_object": None,
"org_object": None,
"token": jwt_token,
"team_id": None,
"user_id": "validated-user",
"user_email": None,
"end_user_id": None,
"org_id": None,
"team_membership": None,
"jwt_claims": {"sub": "user1"},
}
mock_request = MagicMock()
mock_request.url.path = "/v1/chat/completions"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
with (
patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
MagicMock(post_call_failure_hook=AsyncMock(return_value=None)),
),
patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler),
patch(
"litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key",
new_callable=AsyncMock,
return_value=_PendingAutoRegister(
claim_field="sub",
claim_value="user1",
cache_key="jwt_key_mapping:sub:user1",
),
),
patch(
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
new_callable=AsyncMock,
return_value=mock_jwt_result,
),
patch(
"litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping",
new_callable=AsyncMock,
return_value=reused_key,
),
):
call = _user_api_key_auth_builder(
request=mock_request,
api_key=jwt_token,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={"model": "gpt-4o-mini"},
)
if expect_denied:
with pytest.raises(ProxyException, match="not available for this API key"):
await call
return
result = await call
assert result.api_key == "hashed-existing-key"
assert result.user_id == "validated-user"
@pytest.mark.asyncio
@pytest.mark.parametrize("active", [True, False])
async def test_auto_register_first_request_propagates_user_email(active: bool) -> None: