diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 436ceb447a2..d6ed64c2d14 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -401,9 +401,39 @@ def _team_key_generation_check( access_group_ids=data.access_group_ids, ) + _team_key_model_max_budget_check( + team_table=team_table, + user_api_key_dict=user_api_key_dict, + model_max_budget=data.model_max_budget, + ) + return True +def _team_key_model_max_budget_check( + team_table: LiteLLM_TeamTableCachedObj, + user_api_key_dict: UserAPIKeyAuth, + model_max_budget: Mapping[str, Mapping[str, str | float]] | None, +) -> None: + """ + A key-level model_max_budget entry overrides the team's per-model cap for + that key, so a regular member attaching one to a self-serve team key would + opt themselves out of a ceiling a proxy admin imposed. Only a proxy admin + or this team's admin may set model_max_budget on a team key. + """ + if not model_max_budget: + return + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + team_member_object = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) + if team_member_object is not None and team_member_object.role == "admin": + return + raise HTTPException( + status_code=403, + detail=f"Only a proxy admin or team admin can set model_max_budget on a team key. team_id={team_table.team_id}", + ) + + def _personal_key_membership_check( user_api_key_dict: UserAPIKeyAuth, personal_key_generation: PersonalUIKeyGenerationConfig | None, @@ -2362,6 +2392,21 @@ async def _validate_mcp_servers_for_key_update( return normalized_object_permission +def _is_model_max_budget_change( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, +) -> bool: + """ + True when the request would change the key's model_max_budget. Key-level + model_max_budget overrides the team's per-model caps, so changing it gates on + the same admin check as the other budget fields; comparing values (not mere + presence) keeps UI edit flows that re-send the unchanged mapping working. + """ + if "model_max_budget" not in data.model_fields_set: + return False + return (data.model_max_budget or None) != (existing_key_row.model_max_budget or None) + + async def _validate_update_key_data( data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken, @@ -2448,6 +2493,7 @@ async def _validate_update_key_data( (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) or data.spend is not None or "budget_limits" in data.model_fields_set + or _is_model_max_budget_change(data=data, existing_key_row=existing_key_row) ) _existing_metadata = getattr(existing_key_row, "metadata", None) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 315a667c77b..cc9f120216e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1027,6 +1027,56 @@ def _check_team_budget_update_authority( ) +def _check_team_model_max_budget_update_authority( + data: UpdateTeamRequest, + user_api_key_dict: UserAPIKeyAuth, + existing_model_max_budget: Mapping[str, Mapping[str, str | float]] | None, +) -> None: + """ + Restrict who can loosen a team's per-model caps on /team/update. + + Mirrors _check_team_budget_update_authority: a team admin may add new model + caps or lower existing ones, but only a proxy admin may raise a cap, change + its window, or remove an entry (including clearing the whole mapping), since + that undoes a ceiling a proxy admin imposed. + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if "model_max_budget" not in (getattr(data, "model_fields_set", None) or frozenset()): + return + if not existing_model_max_budget: + return + + from types import MappingProxyType + + from litellm.types.utils import BudgetConfig + + existing_configs = MappingProxyType( + {_model: BudgetConfig(**_info) for _model, _info in existing_model_max_budget.items()} + ) + requested_items = data.model_max_budget.items() if data.model_max_budget else () + requested_configs = MappingProxyType({_model: BudgetConfig(**_info) for _model, _info in requested_items}) + for model_name, existing_config in existing_configs.items(): + requested_config = requested_configs.get(model_name) + if requested_config is None: + raise HTTPException( + status_code=403, + detail=f"Only a proxy admin can remove a team's model_max_budget entry. Entry for model={model_name} is missing from the request.", + ) + if requested_config.budget_duration != existing_config.budget_duration: + raise HTTPException( + status_code=403, + detail=f"Only a proxy admin can change a team model_max_budget window. model={model_name} current={existing_config.budget_duration}, requested={requested_config.budget_duration}.", + ) + if existing_config.max_budget is not None and ( + requested_config.max_budget is None or requested_config.max_budget > existing_config.max_budget + ): + raise HTTPException( + status_code=403, + detail=f"Only a proxy admin can raise a team's model_max_budget. model={model_name} current={existing_config.max_budget}, requested={requested_config.max_budget}.", + ) + + def _should_auto_add_team_creator( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], @@ -1965,6 +2015,12 @@ async def update_team( existing_team_max_budget=existing_team_row.max_budget, ) + _check_team_model_max_budget_update_authority( + data=data, + user_api_key_dict=user_api_key_dict, + existing_model_max_budget=existing_team_row.model_max_budget, + ) + if data.model_max_budget is not None: from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_model_max_budget, 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 867ef759fb3..ccd72211c2a 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 @@ -15256,3 +15256,153 @@ async def test_rotate_master_key_rotates_sso_identity_assertions( prisma_client=mock_prisma_client, new_master_key="sk-new-master-key", ) + + +class TestTeamKeyModelMaxBudgetGate: + """Only a proxy admin or the team's admin may attach a key-level + model_max_budget to a team key, since it overrides the team's per-model cap.""" + + def _team(self): + from litellm.proxy._types import Member + + return LiteLLM_TeamTableCachedObj( + team_id="team-1", + members_with_roles=[ + Member(role="user", user_id="member-1"), + Member(role="admin", user_id="team-admin-1"), + ], + ) + + def _mmb(self): + return {"gpt-4o": {"budget_limit": 1000000.0, "time_period": "1d"}} + + def test_regular_member_blocked(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _team_key_model_max_budget_check, + ) + + with pytest.raises(HTTPException) as exc_info: + _team_key_model_max_budget_check( + team_table=self._team(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="member-1" + ), + model_max_budget=self._mmb(), + ) + assert exc_info.value.status_code == 403 + + def test_team_admin_allowed(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _team_key_model_max_budget_check, + ) + + _team_key_model_max_budget_check( + team_table=self._team(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin-1" + ), + model_max_budget=self._mmb(), + ) + + def test_proxy_admin_allowed(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _team_key_model_max_budget_check, + ) + + _team_key_model_max_budget_check( + team_table=self._team(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="root" + ), + model_max_budget=self._mmb(), + ) + + def test_member_without_model_max_budget_allowed(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _team_key_model_max_budget_check, + ) + + _team_key_model_max_budget_check( + team_table=self._team(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="member-1" + ), + model_max_budget=None, + ) + + +class TestIsModelMaxBudgetChange: + """model_max_budget updates gate on the admin check only when the VALUE + changes, so UI flows that re-send the unchanged mapping keep working.""" + + def _row(self, mmb): + return LiteLLM_VerificationToken(token="hashed", model_max_budget=mmb) + + def test_field_absent_is_not_a_change(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_model_max_budget_change, + ) + + assert ( + _is_model_max_budget_change( + data=UpdateKeyRequest(key="sk-1"), + existing_key_row=self._row({"gpt-4o": {"budget_limit": 5.0}}), + ) + is False + ) + + def test_same_value_is_not_a_change(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_model_max_budget_change, + ) + + mmb = {"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}} + assert ( + _is_model_max_budget_change( + data=UpdateKeyRequest(key="sk-1", model_max_budget=mmb), + existing_key_row=self._row(mmb), + ) + is False + ) + + def test_new_value_is_a_change(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_model_max_budget_change, + ) + + assert ( + _is_model_max_budget_change( + data=UpdateKeyRequest( + key="sk-1", + model_max_budget={"gpt-4o": {"budget_limit": 999999.0, "time_period": "1d"}}, + ), + existing_key_row=self._row({"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}}), + ) + is True + ) + + def test_clearing_existing_value_is_a_change(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_model_max_budget_change, + ) + + assert ( + _is_model_max_budget_change( + data=UpdateKeyRequest(key="sk-1", model_max_budget={}), + existing_key_row=self._row({"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}}), + ) + is True + ) + + def test_empty_to_empty_is_not_a_change(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_model_max_budget_change, + ) + + assert ( + _is_model_max_budget_change( + data=UpdateKeyRequest(key="sk-1", model_max_budget={}), + existing_key_row=self._row({}), + ) + is False + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0de2c1ac71d..dad650bfd67 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -10617,3 +10617,128 @@ def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): assert f"u{_MAX_REPORTED_UNKNOWN_USER_IDS}" not in detail assert f"and {500 - _MAX_REPORTED_UNKNOWN_USER_IDS} more" in detail assert len(detail) < 1000 + + +class TestTeamModelMaxBudgetUpdateAuthority: + """Only a proxy admin may loosen team-level per-model caps on /team/update; + a team admin may only add caps or make existing ones stricter.""" + + def _existing(self): + return {"gpt-4o": {"budget_limit": 100.0, "time_period": "1d"}} + + def _team_admin(self): + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin-1") + + def test_proxy_admin_can_do_anything(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget={}), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + existing_model_max_budget=self._existing(), + ) + + def test_non_proxy_admin_cannot_clear_caps(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + with pytest.raises(HTTPException) as exc_info: + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget={}), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=self._existing(), + ) + assert exc_info.value.status_code == 403 + + def test_non_proxy_admin_cannot_remove_one_entry(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + with pytest.raises(HTTPException) as exc_info: + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", + model_max_budget={"claude-3": {"budget_limit": 10.0, "time_period": "1d"}}, + ), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=self._existing(), + ) + assert exc_info.value.status_code == 403 + + def test_non_proxy_admin_cannot_raise_cap(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + with pytest.raises(HTTPException) as exc_info: + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", + model_max_budget={"gpt-4o": {"budget_limit": 1000000.0, "time_period": "1d"}}, + ), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=self._existing(), + ) + assert exc_info.value.status_code == 403 + + def test_non_proxy_admin_cannot_change_window(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + with pytest.raises(HTTPException) as exc_info: + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", + model_max_budget={"gpt-4o": {"budget_limit": 100.0, "time_period": "30d"}}, + ), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=self._existing(), + ) + assert exc_info.value.status_code == 403 + + def test_non_proxy_admin_can_lower_cap_and_add_models(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", + model_max_budget={ + "gpt-4o": {"budget_limit": 50.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 10.0, "time_period": "7d"}, + }, + ), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=self._existing(), + ) + + def test_no_existing_caps_allows_setting(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", + model_max_budget={"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}}, + ), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=None, + ) + + def test_field_not_in_request_is_ignored(self): + from litellm.proxy.management_endpoints.team_endpoints import ( + _check_team_model_max_budget_update_authority, + ) + + _check_team_model_max_budget_update_authority( + data=UpdateTeamRequest(team_id="t1"), + user_api_key_dict=self._team_admin(), + existing_model_max_budget=self._existing(), + )