scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-20 05:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 155

fix(keys): use user-account budget as delegation ceiling for UI session personal keys

Browse source
This commit is contained in:
Devin AI 2026-07-14 17:16:28 +00:00
parent 93b5ca9612
commit 70a16ed9bd
2 changed files with 149 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -49,6 +49,7 @@ from litellm.proxy.auth.auth_checks import (
get_org_object,
get_project_object,
get_team_object,
get_user_object,
)
from litellm.proxy.auth.auth_utils import abbreviate_api_key
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -646,6 +647,43 @@ def _check_budget_limits_delegation_ceiling(
)
async def _resolve_delegation_ceiling(
user_api_key_dict: UserAPIKeyAuth,
team_table: LiteLLM_TeamTableCachedObj | None,
is_ui_session_token: bool,
) -> float | None:
"""
Budget a non-admin caller is allowed to delegate to a newly created key.
A UI session token (team_id == UI_SESSION_TOKEN_TEAM_ID) carries
max_budget == max_ui_session_budget, a per-session chat spend cap (default
$0.25) rather than the caller's real authority, so its ceiling is the
caller's user-account budget resolved from the DB (None once the user is
resolved means the account is uncapped). If the user cannot be resolved the
ceiling fails closed to the session cap. Every other caller delegates from
its own max_budget, falling back to the team budget for a CLI session token
creating a team key.
"""
if is_ui_session_token:
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
user_object = await get_user_object(
user_id=user_api_key_dict.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
except Exception: # noqa: BLE001 # get_user_object raises a bare Exception when the user is absent; fail closed to the session cap
user_object = None
return user_object.max_budget if user_object is not None else user_api_key_dict.max_budget
if user_api_key_dict.max_budget is not None:
return user_api_key_dict.max_budget
if user_api_key_dict.is_session_token and team_table is not None:
return team_table.max_budget
return None
async def validate_team_id_used_in_service_account_request(
team_id: Optional[str],
prisma_client: Optional[PrismaClient],
@ -806,7 +844,8 @@ async def _common_key_generation_helper(
# Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller
# cannot grant a key a higher budget than their own authority.
is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None
is_ui_session_token = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID
is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None
# Session tokens (lite login) carry max_budget=None to avoid a per-session
# LLM spend cap, but that None must not be read as "unlimited delegation
# authority". A personal key (no team) has no team-budget enforcement at
@ -827,10 +866,10 @@ async def _common_key_generation_helper(
)
},
)
delegation_ceiling = (
user_api_key_dict.max_budget
if user_api_key_dict.max_budget is not None
else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None)
delegation_ceiling = await _resolve_delegation_ceiling(
user_api_key_dict=user_api_key_dict,
team_table=team_table,
is_ui_session_token=is_ui_session_token,
)
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value

114
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -13082,14 +13082,14 @@ async def test_ghsa_q775_ui_session_token_team_key_exempt_from_budget_ceiling():
@pytest.mark.asyncio
async def test_ghsa_q775_ui_session_token_personal_key_still_capped():
async def test_ghsa_q775_ui_session_token_personal_key_capped_at_user_account_budget():
"""
Security regression for GHSA-q775: the session-token exemption must NOT extend
to personal keys. A UI/CLI session token (team_id=litellm-dashboard) creating a
key with no data.team_id is still bound by the ceiling; otherwise a session
token - or a leaked one, whose blast radius is the $0.25 chat cap - could mint
an arbitrary-budget personal key, the exact escalation GHSA-q775 closed. Unlike
a team key, nothing else bounds a personal key's spend.
Security regression for GHSA-q775: a UI session token creating a personal key
is bound by the caller's USER-ACCOUNT budget, not the $0.25 session chat cap
(max_ui_session_budget) it carries. Here the user account is capped at $100 and
the request asks for $500, so it must be rejected. This keeps a session token -
or a leaked one - from minting a personal key above the user's real authority
while still letting legitimate within-budget requests through (see #33212).
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
@ -13108,6 +13108,10 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_user_object",
AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1", max_budget=100.0)),
),
):
with pytest.raises((HTTPException, ProxyException)) as exc_info:
await generate_key_fn(
@ -13120,6 +13124,94 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped():
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
assert str(code) == "400"
assert "cannot exceed" in msg.lower()
assert "100" in msg
@pytest.mark.asyncio
async def test_ui_session_token_personal_key_within_user_account_budget_allowed():
"""
Regression for #33212: a non-admin user creating a personal key through the UI
with max_budget below their user-account budget must succeed. The UI session
token carries max_ui_session_budget ($0.25) as a per-session chat cap, which
must not be treated as the delegation ceiling; the user account ($100) is.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
data = GenerateKeyRequest(max_budget=4)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-ui-session",
user_id="user-1",
team_id=UI_SESSION_TOKEN_TEAM_ID,
max_budget=0.25,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_user_object",
AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1", max_budget=100.0)),
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
),
):
result = await _common_key_generation_helper(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
team_table=None,
)
assert result is not None
@pytest.mark.asyncio
async def test_ui_session_token_personal_key_fails_closed_when_user_unresolved():
"""
If the caller's user account cannot be resolved, the ceiling for a UI session
token falls back to the session chat cap ($0.25) so an above-cap personal key
request is still rejected rather than silently allowed.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
data = GenerateKeyRequest(max_budget=500)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-ui-session",
user_id="user-1",
team_id=UI_SESSION_TOKEN_TEAM_ID,
max_budget=0.25,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_user_object",
AsyncMock(side_effect=Exception("user not found")),
),
):
with pytest.raises((HTTPException, ProxyException)) as exc_info:
await _common_key_generation_helper(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
team_table=None,
)
err = exc_info.value
code = getattr(err, "status_code", None) or getattr(err, "code", None)
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
assert str(code) == "400"
assert "cannot exceed" in msg.lower()
@pytest.mark.asyncio
@ -13130,8 +13222,8 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(
With default_key_generate_params.team_id set, a UI session token's personal-key
request (no team_id) would otherwise have team_id auto-filled before the ceiling
check, flipping is_ui_session_team_key to True and bypassing the ceiling. The
request must still be rejected. Mirrors how _requested_max_budget is captured
before defaults run.
request (above the user-account budget) must still be rejected. Mirrors how
_requested_max_budget is captured before defaults run.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
@ -13152,6 +13244,10 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
patch("litellm.default_key_generate_params", {"team_id": "injected-team"}),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_user_object",
AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1", max_budget=100.0)),
),
):
with pytest.raises((HTTPException, ProxyException)) as exc_info:
await _common_key_generation_helper(