From 8c79bbc1eb51c9389f501eec3214961d22cf32bb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:47:47 -0400 Subject: [PATCH] fix(proxy): use user budget for dashboard key creation --- litellm/proxy/auth/auth_checks.py | 1 + .../key_management_endpoints.py | 18 +++-- litellm/proxy/management_endpoints/ui_sso.py | 2 +- .../proxy/auth/test_auth_checks.py | 2 + .../test_key_management_endpoints.py | 68 +++++++++++++++++++ 5 files changed, 84 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3d8fed18423..11110568ac2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2653,6 +2653,7 @@ class ExperimentalUIJWTToken: models=user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), + user_max_budget=user_info.max_budget, ) return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e190e8b19d..f1702f35834 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -921,7 +921,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 @@ -942,11 +943,16 @@ async def _common_key_generation_helper( ) }, ) - delegation_ceiling: Final = ( - 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) - ) + # The dashboard key's max_budget only caps playground spend. Its owner's + # personal budget is the authority for creating personal keys. + if is_ui_session_token and user_api_key_dict.user_max_budget is not None: + delegation_ceiling = user_api_key_dict.user_max_budget + elif user_api_key_dict.max_budget is not None: + delegation_ceiling = user_api_key_dict.max_budget + else: + delegation_ceiling = ( + team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None + ) if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value and not is_ui_session_team_key diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index b87ad8597dc..6cd8e25d755 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3592,7 +3592,7 @@ class SSOAuthenticationHandler: user_id=user_defined_values["user_id"], user_role=user_defined_values["user_role"] or user_role, models=[], - max_budget=litellm.max_ui_session_budget, + max_budget=user_defined_values["max_budget"], ) if _user_info is None: raise HTTPException( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 28eda6633e8..3ffade7df6e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -128,6 +128,7 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value assert token_data["models"] == ["gpt-3.5-turbo"] assert token_data["max_budget"] == litellm.max_ui_session_budget + assert token_data["user_max_budget"] == 100.0 # Verify expiration time is set and valid (Experimental UI uses fixed 10-min expiry) assert "expires" in token_data @@ -276,6 +277,7 @@ def test_get_key_object_from_ui_hash_key_valid( assert key_object.user_role == LitellmUserRoles.PROXY_ADMIN assert key_object.models == ["gpt-3.5-turbo"] assert key_object.max_budget == litellm.max_ui_session_budget + assert key_object.user_max_budget == 100.0 def test_get_key_object_from_ui_hash_key_invalid(): 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 bdf09a95e4b..1eb444db670 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 @@ -13362,6 +13362,74 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped(): assert "cannot exceed" in msg.lower() +@pytest.mark.asyncio +async def test_ui_session_personal_key_uses_user_budget_as_delegation_ceiling(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=20) + 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=1.0, + user_max_budget=30.0, + ) + generated_key = {"key": "sk-generated"} + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + 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.generate_key_helper_fn", + new=AsyncMock(return_value=generated_key), + ) as mock_generate_key, + ): + 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.key == "sk-generated" + assert mock_generate_key.await_args.kwargs["key_max_budget"] == 20 + + +@pytest.mark.asyncio +async def test_ui_session_personal_key_cannot_exceed_user_budget(): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + data = GenerateKeyRequest(max_budget=31) + 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=1.0, + user_max_budget=30.0, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + 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"), + pytest.raises(HTTPException) 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, + ) + + assert exc_info.value.status_code == 400 + assert "own max_budget (30.0)" in str(exc_info.value.detail) + + @pytest.mark.asyncio async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(): """