diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 576585ee9a3..0dd7f1c8271 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3370,6 +3370,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 95ccb7bbe0b..328b191bc46 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1068,7 +1068,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 @@ -1089,11 +1090,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 1ba90725eff..b8a2363ef36 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3646,7 +3646,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 ea2c925212d..565903fba36 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -142,6 +142,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 @@ -270,6 +271,7 @@ def test_get_key_object_from_ui_hash_key_valid(valid_sso_user_defined_values, mo 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 2ac52da57df..4f730a66152 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 @@ -14101,6 +14101,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(): """