Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
devin-ai-integration wants to merge 1 commit into
base: litellm_internal_staging
Choose a base branch
from
+149 −14
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 44 additions & 5 deletions litellm/proxy/management_endpoints/key_management_endpoints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
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
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
114 changes: 105 additions & 9 deletions tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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(
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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(
Expand Down
Loading