From 70a16ed9bd3d34e5ef2cd5090f15f06f1fe19a09 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:16:28 +0000 Subject: [PATCH] fix(keys): use user-account budget as delegation ceiling for UI session personal keys --- .../key_management_endpoints.py | 49 +++++++- .../test_key_management_endpoints.py | 114 ++++++++++++++++-- 2 files changed, 149 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a3679b84bd68..2458a683024b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index c9abdf09a5dd..5b941c3dc33c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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(