From 40e4ade6a6d5824a072e90a391cf8b2d32675a48 Mon Sep 17 00:00:00 2001 From: Justin Israel Date: Tue, 20 Oct 2026 07:40:38 +1300 Subject: [PATCH] feat(proxy): opt-in self-serve budget edits for key owners on /key/update (#44763) * feat(proxy): opt-in self-serve budget edits for key owners on /key/update Adds general_settings.self_serve_budget_policy (disabled by default) so a non-admin can change max_budget or budget_limits on a key they created and still own. lower_only allows adding or lowering limits. ceiling also allows raising or removing them when the owner could create that key with the new budget via /key/generate (same member permissions and delegation ceiling). soft_budget and spend stay admin-only * fix(proxy): sync dashboard API types and cover the UI session cap on update * test(proxy): cover team role restriction for owner budget raises Drop two unreachable branches from the non-raising key generation check. * fix(proxy): close self-serve budget bypasses for linked budgets, null caps and retained limits Keys with a linked budget_id need an admin to change max_budget, owners cannot leave a key uncapped when an upperbound max_budget is configured, and the ceiling check now runs on the budget the key would carry after the update. * fix(proxy): keep open key budget windows and the admin path on self-serve budget edits A resubmitted budget_limits window now keeps its stored reset_at while that period is still running, so adding or editing one window no longer restarts the others. Team and org admins answer False on a ceiling loosen so they fall through to _check_key_admin_access as before. The self-serve tests feed the team through the cache and prisma mocks instead of patching get_team_object, and every helper is typed * perf(proxy): group key budget windows by duration in one pass The per-duration minimum used to rescan the whole window list once per distinct duration, so a request with thousands of windows pinned the event loop before any permission check ran. Sort once and take each group's minimum instead * refactor(proxy): take budget windows as a Sequence in the delegation ceiling checks * test(proxy): assert allowed self-serve budget edits return no denial Also drop the unreachable assert_never arm from the exhaustive budget change match. --------- Co-authored-by: ryan-crabbe-berri --- litellm/proxy/_types.py | 4 + .../management_endpoints/key_budget_change.py | 135 +++++ .../key_management_endpoints.py | 280 ++++++++--- .../team_member_permission_checks.py | 17 +- .../test_key_budget_change.py | 467 ++++++++++++++++++ .../test_key_management_endpoints.py | 45 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 7 files changed, 874 insertions(+), 79 deletions(-) create mode 100644 litellm/proxy/management_endpoints/key_budget_change.py create mode 100644 tests/unit/proxy/management_endpoints/test_key_budget_change.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bc32978d477..e4f77ac0005 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2793,6 +2793,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools', 'config_overrides'. If not set, all objects are loaded (default behavior).", ) + self_serve_budget_policy: Literal["disabled", "lower_only", "ceiling"] | None = Field( + None, + description="Lets a non-admin change max_budget or budget_limits on a key they created and still own via /key/update. 'disabled' (default) requires an admin for any budget change. 'lower_only' lets the owner add or lower limits. 'ceiling' also lets the owner raise or remove limits when they could create that key with the new budget via /key/generate.", + ) user_mcp_management_mode: UserMCPManagementMode | None = Field( None, description="Controls how non-admin users interact with MCP servers in the dashboard. 'restricted' shows only accessible servers, 'view_all' lists every server in read-only mode.", diff --git a/litellm/proxy/management_endpoints/key_budget_change.py b/litellm/proxy/management_endpoints/key_budget_change.py new file mode 100644 index 00000000000..3143ef53d76 --- /dev/null +++ b/litellm/proxy/management_endpoints/key_budget_change.py @@ -0,0 +1,135 @@ +"""Classifies a /key/update request's budget edits as tightening or loosening the key's stored limits.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import groupby +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.models.team import BudgetLimitEntry +from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy._types import UpdateKeyRequest + + +@dataclass(frozen=True, slots=True) +class KeyBudgetUnchanged: + kind: Literal["unchanged"] = "unchanged" + + +@dataclass(frozen=True, slots=True) +class KeyBudgetTightened: + kind: Literal["tightened"] = "tightened" + + +@dataclass(frozen=True, slots=True) +class KeyBudgetLoosened: + kind: Literal["loosened"] = "loosened" + + +@dataclass(frozen=True, slots=True) +class KeyBudgetAdminOnly: + kind: Literal["admin_only"] = "admin_only" + + +KeyBudgetChange: TypeAlias = KeyBudgetUnchanged | KeyBudgetTightened | KeyBudgetLoosened | KeyBudgetAdminOnly + +_FieldChange: TypeAlias = Literal["unchanged", "tightened", "loosened"] + +SELF_SERVE_BUDGET_POLICY_SETTING: Final = "self_serve_budget_policy" +SelfServeBudgetPolicy: TypeAlias = Literal["disabled", "lower_only", "ceiling"] + +_STORED_WINDOWS: Final = TypeAdapter(list[BudgetLimitEntry]) +_POLICY: Final = TypeAdapter(SelfServeBudgetPolicy) + + +def resolve_self_serve_budget_policy(general_settings: Mapping[str, object]) -> SelfServeBudgetPolicy: + configured: Final = general_settings.get(SELF_SERVE_BUDGET_POLICY_SETTING, "disabled") + try: + return _POLICY.validate_python(configured) + except ValidationError: + verbose_proxy_logger.warning( + "general_settings.%s=%r is not one of disabled, lower_only, ceiling; using disabled", + SELF_SERVE_BUDGET_POLICY_SETTING, + configured, + ) + return "disabled" + + +def _max_budget_change(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> _FieldChange: + if "max_budget" not in data.model_fields_set or data.max_budget == existing.max_budget: + return "unchanged" + if data.max_budget is None: + return "loosened" + if existing.max_budget is None or data.max_budget < existing.max_budget: + return "tightened" + return "loosened" + + +def _duration(window: BudgetLimitEntry) -> str: + return window.budget_duration + + +def _caps_by_duration(windows: Sequence[BudgetLimitEntry]) -> Mapping[str, float]: + grouped: Final = groupby(sorted(windows, key=_duration), key=_duration) + return MappingProxyType({duration: min(w.max_budget for w in group) for duration, group in grouped}) + + +def _stored_caps(existing: LiteLLM_VerificationToken) -> Mapping[str, float] | None: + try: + return _caps_by_duration(_STORED_WINDOWS.validate_python(existing.budget_limits or [])) + except ValidationError: + return None + + +def _budget_limits_change(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> _FieldChange: + """ + Window spend counters are keyed by the exact `budget_duration` string, so a window only + counts as kept when the same string survives with a cap no higher than before. + """ + if "budget_limits" not in data.model_fields_set: + return "unchanged" + stored: Final = _stored_caps(existing) + if stored is None: + return "loosened" + requested: Final = _caps_by_duration(data.budget_limits or ()) + kept_or_lowered: Final = all(d in requested and requested[d] <= cap for d, cap in stored.items()) + return "tightened" if kept_or_lowered else "loosened" + + +@dataclass(frozen=True, slots=True) +class EffectiveKeyBudget: + max_budget: float | None + budget_limits: tuple[BudgetLimitEntry, ...] + + +def stored_budget_windows(existing: LiteLLM_VerificationToken) -> tuple[BudgetLimitEntry, ...]: + """The key's stored windows, or none when the column holds something the window model cannot read.""" + try: + return tuple(_STORED_WINDOWS.validate_python(existing.budget_limits or [])) + except ValidationError: + return () + + +def effective_key_budget(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> EffectiveKeyBudget: + """The budget the key would carry after the update, with omitted fields keeping their stored values.""" + max_budget: Final = data.max_budget if "max_budget" in data.model_fields_set else existing.max_budget + if "budget_limits" in data.model_fields_set: + return EffectiveKeyBudget(max_budget=max_budget, budget_limits=tuple(data.budget_limits or ())) + return EffectiveKeyBudget(max_budget=max_budget, budget_limits=stored_budget_windows(existing)) + + +def classify_key_budget_change(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> KeyBudgetChange: + overrides_linked_budget: Final = ( + existing.budget_id is not None and _max_budget_change(data, existing) != "unchanged" + ) + if data.spend is not None or "soft_budget" in data.model_fields_set or overrides_linked_budget: + return KeyBudgetAdminOnly() + changes: Final = frozenset((_max_budget_change(data, existing), _budget_limits_change(data, existing))) + if "loosened" in changes: + return KeyBudgetLoosened() + if "tightened" in changes: + return KeyBudgetTightened() + return KeyBudgetUnchanged() diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e87ca70035c..022eba499a4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -106,6 +106,17 @@ from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # l validate_budget_duration, validate_finite_spend, ) +from litellm.proxy.management_endpoints.key_budget_change import ( + KeyBudgetAdminOnly, + KeyBudgetChange, + KeyBudgetLoosened, + KeyBudgetTightened, + KeyBudgetUnchanged, + classify_key_budget_change, + effective_key_budget, + resolve_self_serve_budget_policy, + stored_budget_windows, +) from litellm.proxy.management_endpoints.model_management_endpoints import ( # noqa: F401 # legacy module exports _add_model_to_db, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export add_model_to_db, @@ -190,6 +201,7 @@ from litellm.types.router import Deployment from litellm.types.utils import ( BudgetConfig, PersonalUIKeyGenerationConfig, + StandardKeyGenerationConfig, TeamUIKeyGenerationConfig, ) @@ -848,6 +860,33 @@ def key_generation_check( return _personal_key_generation_check(user_api_key_dict=user_api_key_dict, data=data) +def _caller_may_generate_key( + team_table: LiteLLM_TeamTableCachedObj | None, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: + """Non-raising mirror of the role and member-permission gates in `key_generation_check`.""" + settings: Final = litellm.key_generation_settings or StandardKeyGenerationConfig() + if team_table is None: + personal: Final = settings.get("personal_key_generation") + return ( + personal is None + or "allowed_user_roles" not in personal + or user_api_key_dict.user_role in personal["allowed_user_roles"] + ) + team_settings: Final = settings.get( + "team_key_generation", TeamUIKeyGenerationConfig(allowed_team_member_roles=["admin", "user"]) + ) + caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict) + if caller_team_role is None or ( + "allowed_team_member_roles" in team_settings + and caller_team_role not in team_settings["allowed_team_member_roles"] + ): + return False + return caller_team_role == "admin" or TeamMemberPermissionChecks.is_route_granted_to_team_members( + team_table=team_table, route=KeyManagementRoutes.KEY_GENERATE.value + ) + + def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) -> None: """Key-level logging writes go through key metadata, not /team/callback. @@ -1072,7 +1111,7 @@ _check_permissions_caller_permission: Final = check_permissions_caller_permissio def _check_budget_limits_delegation_ceiling( - budget_limits: list[BudgetLimitEntry] | None, + budget_limits: Sequence[BudgetLimitEntry] | None, delegation_ceiling: float | None, user_api_key_dict: UserAPIKeyAuth, is_ui_session_team_key: bool, @@ -1122,6 +1161,67 @@ def _check_budget_limits_delegation_ceiling( ) +def _enforce_budget_delegation_ceiling( + requested_max_budget: float | None, + requested_team_id: str | None, + budget_limits: Sequence[BudgetLimitEntry] | None, + team_table: LiteLLM_TeamTableCachedObj | None, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller + # cannot grant a key a higher budget than their own authority. + # UI session personal keys are capped by user_max_budget when it is available. + is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID + is_ui_session_team_key: Final = is_ui_session_token and requested_team_id is not None + if ( + user_api_key_dict.is_session_token + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + and not is_ui_session_team_key + and requested_max_budget is not None + and team_table is None + ): + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"max_budget ({requested_max_budget}) cannot be set without " + "specifying team_id when using a CLI session token." + ) + }, + ) + delegation_ceiling: Final = ( + user_api_key_dict.user_max_budget + if is_ui_session_token and user_api_key_dict.user_max_budget is not None + else 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) + ) + if ( + user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + and not is_ui_session_team_key + and requested_max_budget is not None + and delegation_ceiling is not None + and requested_max_budget > delegation_ceiling + ): + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"max_budget ({requested_max_budget}) cannot exceed the caller's " + f"own max_budget ({delegation_ceiling})." + ) + }, + ) + + _check_budget_limits_delegation_ceiling( + budget_limits=budget_limits, + delegation_ceiling=delegation_ceiling, + user_api_key_dict=user_api_key_dict, + is_ui_session_team_key=is_ui_session_team_key, + team_table=team_table, + ) + + async def validate_team_id_used_in_service_account_request( team_id: str | None, prisma_client: PrismaClient | None, @@ -1364,57 +1464,12 @@ async def _common_key_generation_helper( # check if user set upperbound key/generate params on config.yaml _enforce_upperbound_key_params(data, fill_defaults=True) - # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller - # cannot grant a key a higher budget than their own authority. - # UI session personal keys are capped by user_max_budget when it is available. - is_ui_session_token: Final = 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 - if ( - user_api_key_dict.is_session_token - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not is_ui_session_team_key - and _requested_max_budget is not None - and team_table is None - ): - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"max_budget ({_requested_max_budget}) cannot be set without " - "specifying team_id when using a CLI session token." - ) - }, - ) - delegation_ceiling: Final = ( - user_api_key_dict.user_max_budget - if is_ui_session_token and user_api_key_dict.user_max_budget is not None - else 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) - ) - if ( - user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not is_ui_session_team_key - and _requested_max_budget is not None - and delegation_ceiling is not None - and _requested_max_budget > delegation_ceiling - ): - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"max_budget ({_requested_max_budget}) cannot exceed the caller's " - f"own max_budget ({delegation_ceiling})." - ) - }, - ) - - _check_budget_limits_delegation_ceiling( + _enforce_budget_delegation_ceiling( + requested_max_budget=_requested_max_budget, + requested_team_id=_requested_team_id, budget_limits=data.budget_limits, - delegation_ceiling=delegation_ceiling, - user_api_key_dict=user_api_key_dict, - is_ui_session_team_key=is_ui_session_team_key, team_table=team_table, + user_api_key_dict=user_api_key_dict, ) check_permissions_caller_permission( data=data, @@ -2750,16 +2805,12 @@ async def prepare_key_update_data( non_default_values["budget_duration"] = budget_duration if "budget_limits" in non_default_values: - raw_windows: Final = non_default_values["budget_limits"] - if raw_windows: - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - initialized_windows: Final = [] - for window in raw_windows: - w = window if isinstance(window, dict) else window.model_dump() - w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat() - initialized_windows.append(w) - non_default_values["budget_limits"] = json.dumps(initialized_windows) + requested_windows: Final = _budget_limit_windows(data.budget_limits) + if requested_windows: + stored_windows: Final = stored_budget_windows(existing_key_row) + non_default_values["budget_limits"] = json.dumps( + [_window_with_reset_at(window, stored_windows) for window in requested_windows] + ) else: # [] / None clears the field; prisma-client-py has no DbNull # sentinel for Json? columns, so store the JSON literal null @@ -3327,6 +3378,62 @@ async def _acting_as_team_admin_for_key_update( return True +async def _owner_may_change_key_budget( + budget_change: KeyBudgetChange, + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> bool: + """ + Whether the key's owner may make this budget change without admin rights, per + general_settings.self_serve_budget_policy. `ceiling` lets the owner loosen only when they + could have created the key with the new budget anyway (same generate permissions and + delegation ceiling as /key/generate). Team and org admins answer False on a loosen so they + keep the admin path, which has no ceiling. + """ + policy: Final = resolve_self_serve_budget_policy(_general_settings()) + if policy == "disabled": + return False + match budget_change: + case KeyBudgetUnchanged() | KeyBudgetTightened(): + return True + case KeyBudgetAdminOnly(): + return False + case KeyBudgetLoosened(): + if policy == "lower_only": + return False + team_id: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id + team_table: Final = ( + None + if team_id is None + else await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + ) + if team_table is not None and await get_team_access().allows( + user_api_key_dict, team_table, TEAM_OR_ORG_ADMIN + ): + return False + if not _caller_may_generate_key(team_table=team_table, user_api_key_dict=user_api_key_dict): + return False + effective: Final = effective_key_budget(data=data, existing=existing_key_row) + upperbound: Final = litellm.upperbound_key_generate_params + if effective.max_budget is None and upperbound is not None and upperbound.max_budget is not None: + return False + _enforce_budget_delegation_ceiling( + requested_max_budget=effective.max_budget, + requested_team_id=team_id, + budget_limits=effective.budget_limits, + team_table=team_table, + user_api_key_dict=user_api_key_dict, + ) + return True + + async def _validate_update_key_data( data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken, @@ -3420,22 +3527,18 @@ async def _validate_update_key_data( # - Anyone else (non-PROXY_ADMIN, not the owner, not a team member # on a team key): must pass _check_key_admin_access (PROXY_ADMIN # / key-owner / team-admin / org-admin of the key). - # - max_budget / spend / budget_limits: always require the admin - # check, even for the key owner or a team member (matches the - # existing admin-only budget semantics). budget_limits uses - # model_fields_set because an explicit null/[] clears the field - # and must gate the same as setting or changing it. + # - max_budget / budget_limits: admin check unless + # general_settings.self_serve_budget_policy lets the key owner make the + # change (_owner_may_change_key_budget). + # - soft_budget: admin check, since it writes the key's budget row, + # which can be a shared /budget/new tier. # - spend gates on presence alone (not a value diff): the DB spend # lags the live cross-pod counter, so letting an "unchanged" spend # through the non-admin path would let a key owner / team member # overwrite the live counter below real usage and silently weaken # enforcement. - _is_budget_change: Final = ( - (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 "soft_budget" in data.model_fields_set - ) + budget_change: Final = classify_key_budget_change(data=data, existing=existing_key_row) + _is_budget_change: Final = not isinstance(budget_change, KeyBudgetUnchanged) _existing_metadata: Final = getattr(existing_key_row, "metadata", None) _existing_throttle: Final = ( @@ -3484,8 +3587,19 @@ async def _validate_update_key_data( # non-budget change means the caller was authorized — skip the redundant # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key: Final = getattr(existing_key_row, "team_id", None) is not None - can_skip_admin_check: Final = (caller_is_creator or _key_is_team_key) and not ( - _is_budget_change or is_project_change + budget_change_allowed: Final = not _is_budget_change or ( + caller_is_creator + and await _owner_may_change_key_budget( + budget_change=budget_change, + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + ) + ) + can_skip_admin_check: Final = ( + (caller_is_creator or _key_is_team_key) and budget_change_allowed and not is_project_change ) if (not _is_proxy_admin) and not can_skip_admin_check: hashed_key: Final = existing_key_row.token @@ -6459,6 +6573,26 @@ def _budget_limit_windows(budget_limits: Sequence[object] | str | None) -> tuple return tuple(raw_window if isinstance(raw_window, dict) else raw_window.model_dump() for raw_window in raw_windows) +def _window_with_reset_at( + window: Mapping[str, object], stored_windows: Sequence[BudgetLimitEntry] +) -> Mapping[str, object]: + """A resubmitted window keeps its running period, so editing one window never restarts the others.""" + duration: Final = str(window["budget_duration"]) + now: Final = datetime.now(timezone.utc) + open_reset_at: Final = next( + (w.reset_at for w in stored_windows if w.budget_duration == duration and _is_open_reset_at(w.reset_at, now)), + None, + ) + reset_at: Final = open_reset_at if open_reset_at is not None else get_budget_reset_time(budget_duration=duration) + return {**window, "reset_at": reset_at.isoformat()} + + +def _is_open_reset_at(reset_at: datetime | None, now: datetime) -> bool: + if reset_at is None: + return False + return (reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc)) > now + + def _advance_one_key_budget_window(window: Mapping[str, object]) -> Mapping[str, object]: """Restart one budget window from now, by advancing its `reset_at`. diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 3e6d0e47265..834b032b04f 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -101,6 +101,13 @@ class TeamMemberPermissionChecks: code=401, ) + @staticmethod + def is_route_granted_to_team_members(team_table: LiteLLM_TeamTableCachedObj, route: str) -> bool: + granted_routes: Final = TeamMemberPermissionChecks._get_list_of_route_enum_as_str( + TeamMemberPermissionChecks.get_permissions_for_team_member(team_table=team_table) + ) + return RouteChecks.check_route_access(route=route, allowed_routes=granted_routes) + @staticmethod def does_team_member_have_permissions_for_endpoint( team_member_role: Literal["admin", "user"] | None, @@ -118,12 +125,10 @@ class TeamMemberPermissionChecks: if team_member_role == "admin": return True - _team_member_permissions: Final = TeamMemberPermissionChecks.get_permissions_for_team_member( - team_table=team_table, - ) - team_member_permissions = TeamMemberPermissionChecks._get_list_of_route_enum_as_str(_team_member_permissions) - - if not RouteChecks.check_route_access(route=route, allowed_routes=team_member_permissions): + if not TeamMemberPermissionChecks.is_route_granted_to_team_members(team_table=team_table, route=route): + team_member_permissions: Final = TeamMemberPermissionChecks._get_list_of_route_enum_as_str( + TeamMemberPermissionChecks.get_permissions_for_team_member(team_table=team_table) + ) raise ProxyException( message=f"Team member does not have permissions for endpoint: {route}. You only have access to the following endpoints: {team_member_permissions} for team {team_table.team_id}. To create keys for this team, please ask your proxy admin to check the team member permission settings and update the settings to allow team member users to create keys.", type=ProxyErrorTypes.team_member_permission_error, diff --git a/tests/unit/proxy/management_endpoints/test_key_budget_change.py b/tests/unit/proxy/management_endpoints/test_key_budget_change.py new file mode 100644 index 00000000000..baf8d870117 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_key_budget_change.py @@ -0,0 +1,467 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + LitellmUserRoles, + Member, + ProxyException, + UpdateKeyRequest, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.management_endpoints.key_budget_change import ( + KeyBudgetAdminOnly, + KeyBudgetChange, + KeyBudgetLoosened, + KeyBudgetTightened, + KeyBudgetUnchanged, + classify_key_budget_change, + resolve_self_serve_budget_policy, +) +from litellm.proxy.management_endpoints.key_management_endpoints import _validate_update_key_data +from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams + +_OWNER: Final = "owner-1" +_TEAM_ADMIN: Final = "team-admin" +_TEAM: Final = "team-1" +_WEEKLY: Final[Mapping[str, object]] = MappingProxyType( + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": "2026-10-08T00:00:00+00:00"} +) + + +def _key( + *, + user_id: str = _OWNER, + created_by: str = _OWNER, + team_id: str | None = None, + max_budget: float | None = None, + budget_limits: Sequence[Mapping[str, object]] | None = None, + budget_id: str | None = None, +) -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token="hashed-owner-key", + user_id=user_id, + created_by=created_by, + team_id=team_id, + max_budget=max_budget, + budget_limits=None if budget_limits is None else [dict(window) for window in budget_limits], + budget_id=budget_id, + metadata={}, + ) + + +def _window(duration: str, cap: float) -> Mapping[str, object]: + return MappingProxyType({"budget_duration": duration, "max_budget": cap}) + + +@pytest.mark.parametrize( + ("update", "stored", "expected"), + [ + (UpdateKeyRequest(key="sk-1", key_alias="renamed"), _key(max_budget=10.0), KeyBudgetUnchanged()), + (UpdateKeyRequest(key="sk-1", max_budget=10.0), _key(max_budget=10.0), KeyBudgetUnchanged()), + (UpdateKeyRequest(key="sk-1", max_budget=5.0), _key(max_budget=10.0), KeyBudgetTightened()), + (UpdateKeyRequest(key="sk-1", max_budget=5.0), _key(), KeyBudgetTightened()), + (UpdateKeyRequest(key="sk-1", max_budget=20.0), _key(max_budget=10.0), KeyBudgetLoosened()), + (UpdateKeyRequest(key="sk-1", max_budget=None), _key(max_budget=10.0), KeyBudgetLoosened()), + (UpdateKeyRequest(key="sk-1", budget_limits=[_window("24h", 5.0)]), _key(), KeyBudgetTightened()), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_WEEKLY, _window("24h", 5.0)]), + _key(budget_limits=[_WEEKLY]), + KeyBudgetTightened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("7d", 50.0)]), + _key(budget_limits=[_WEEKLY]), + KeyBudgetTightened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("7d", 40.0)]), + _key(budget_limits=[_WEEKLY]), + KeyBudgetTightened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("7d", 60.0)]), + _key(budget_limits=[_WEEKLY]), + KeyBudgetLoosened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("24h", 5.0)]), + _key(budget_limits=[_WEEKLY]), + KeyBudgetLoosened(), + ), + (UpdateKeyRequest(key="sk-1", budget_limits=[]), _key(budget_limits=[_WEEKLY]), KeyBudgetLoosened()), + (UpdateKeyRequest(key="sk-1", budget_limits=None), _key(budget_limits=[_WEEKLY]), KeyBudgetLoosened()), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("1d", 5.0)]), + _key(budget_limits=[_window("24h", 5.0)]), + KeyBudgetLoosened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("7d", 40.0)]), + _key(budget_limits=[_WEEKLY, _window("7d", 30.0)]), + KeyBudgetLoosened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("7d", 40.0), _window("24h", 5.0)]), + _key(budget_limits=[_window("7d", 30.0), _window("24h", 5.0), _window("7d", 50.0)]), + KeyBudgetLoosened(), + ), + ( + UpdateKeyRequest(key="sk-1", budget_limits=[_window("24h", 5.0)]), + _key(budget_limits=[{"budget_duration": "24h"}]), + KeyBudgetLoosened(), + ), + ( + UpdateKeyRequest(key="sk-1", max_budget=5.0, budget_limits=[]), + _key(max_budget=10.0, budget_limits=[_WEEKLY]), + KeyBudgetLoosened(), + ), + (UpdateKeyRequest(key="sk-1", spend=0.0), _key(), KeyBudgetAdminOnly()), + (UpdateKeyRequest(key="sk-1", soft_budget=1.0), _key(), KeyBudgetAdminOnly()), + (UpdateKeyRequest(key="sk-1", soft_budget=None, max_budget=5.0), _key(max_budget=10.0), KeyBudgetAdminOnly()), + ], +) +def test_classify_key_budget_change( + update: UpdateKeyRequest, stored: LiteLLM_VerificationToken, expected: KeyBudgetChange +) -> None: + assert classify_key_budget_change(update, stored) == expected + + +@pytest.mark.parametrize( + ("settings", "expected"), + [ + ({}, "disabled"), + ({"self_serve_budget_policy": "lower_only"}, "lower_only"), + ({"self_serve_budget_policy": "ceiling"}, "ceiling"), + ({"self_serve_budget_policy": "everything"}, "disabled"), + ({"self_serve_budget_policy": True}, "disabled"), + ], +) +def test_resolve_self_serve_budget_policy(settings: Mapping[str, object], expected: str) -> None: + assert resolve_self_serve_budget_policy(settings) == expected + + +def _team(member_permissions: Sequence[str]) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj( + team_id=_TEAM, + members_with_roles=[Member(user_id=_TEAM_ADMIN, role="admin"), Member(user_id=_OWNER, role="user")], + team_member_permissions=list(member_permissions), + ) + + +def _caller( + *, + user_id: str = _OWNER, + max_budget: float | None = None, + team_id: str | None = None, + user_max_budget: float | None = None, +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + max_budget=max_budget, + team_id=team_id, + user_max_budget=user_max_budget, + ) + + +async def _update( + monkeypatch: pytest.MonkeyPatch, + *, + policy: str | None, + data: UpdateKeyRequest, + existing: LiteLLM_VerificationToken, + team: LiteLLM_TeamTableCachedObj | None = None, + caller: UserAPIKeyAuth | None = None, +) -> None: + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {} if policy is None else {"self_serve_budget_policy": policy}, + ) + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + prisma.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=caller or _caller(), + llm_router=None, + premium_user=False, + prisma_client=prisma, + user_api_key_cache=UserApiKeyCache(), + ) + + +async def _denial( + monkeypatch: pytest.MonkeyPatch, + *, + policy: str | None, + data: UpdateKeyRequest, + existing: LiteLLM_VerificationToken, + team: LiteLLM_TeamTableCachedObj | None = None, + caller: UserAPIKeyAuth | None = None, +) -> HTTPException | None: + try: + await _update(monkeypatch, policy=policy, data=data, existing=existing, team=team, caller=caller) + except HTTPException as exc: + return exc + return None + + +_ADD_DAILY_WINDOW: Final = UpdateKeyRequest(key="sk-1", budget_limits=[_WEEKLY, _window("24h", 5.0)]) +_RAISE_WEEKLY_WINDOW: Final = UpdateKeyRequest(key="sk-1", budget_limits=[_window("7d", 80.0)]) + + +def _team_key_with_weekly(*, user_id: str = _OWNER, created_by: str = _OWNER) -> LiteLLM_VerificationToken: + return _key(user_id=user_id, created_by=created_by, team_id=_TEAM, budget_limits=[_WEEKLY]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("policy", "data", "member_permissions"), + [ + ("lower_only", _ADD_DAILY_WINDOW, ["/key/update"]), + ("ceiling", _ADD_DAILY_WINDOW, ["/key/update"]), + ("ceiling", _RAISE_WEEKLY_WINDOW, ["/key/update", "/key/generate"]), + ], +) +async def test_owner_changes_team_key_budget_when_policy_allows( + monkeypatch: pytest.MonkeyPatch, policy: str, data: UpdateKeyRequest, member_permissions: Sequence[str] +) -> None: + assert ( + await _denial( + monkeypatch, + policy=policy, + data=data, + existing=_team_key_with_weekly(), + team=_team(member_permissions), + ) + is None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("policy", "data", "member_permissions"), + [ + (None, _ADD_DAILY_WINDOW, ["/key/update", "/key/generate"]), + ("disabled", _ADD_DAILY_WINDOW, ["/key/update", "/key/generate"]), + ("lower_only", _RAISE_WEEKLY_WINDOW, ["/key/update", "/key/generate"]), + ("ceiling", _RAISE_WEEKLY_WINDOW, ["/key/update"]), + ("ceiling", UpdateKeyRequest(key="sk-1", soft_budget=1.0), ["/key/update", "/key/generate"]), + ("ceiling", UpdateKeyRequest(key="sk-1", spend=0.0), ["/key/update", "/key/generate"]), + ], +) +async def test_owner_team_key_budget_change_needs_admin( + monkeypatch: pytest.MonkeyPatch, policy: str | None, data: UpdateKeyRequest, member_permissions: Sequence[str] +) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy=policy, + data=data, + existing=_team_key_with_weekly(), + team=_team(member_permissions), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_key_created_by_someone_else_needs_admin_even_to_tighten(monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=_ADD_DAILY_WINDOW, + existing=_team_key_with_weekly(created_by=_TEAM_ADMIN), + team=_team(["/key/update", "/key/generate"]), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_team_member_without_update_permission_is_still_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.raises(ProxyException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=_ADD_DAILY_WINDOW, + existing=_team_key_with_weekly(), + team=_team(["/key/generate"]), + ) + assert str(exc.value.code) == "401" + + +@pytest.mark.asyncio +async def test_ceiling_lets_owner_raise_personal_key_budget(monkeypatch: pytest.MonkeyPatch) -> None: + assert ( + await _denial( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", max_budget=30.0), + existing=_key(max_budget=10.0), + ) + is None + ) + + +@pytest.mark.asyncio +async def test_ceiling_caps_owner_raise_at_callers_own_budget(monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", max_budget=30.0), + existing=_key(max_budget=10.0), + caller=_caller(max_budget=20.0), + ) + assert exc.value.status_code == 400 + assert "cannot exceed the caller's own max_budget (20.0)" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_ceiling_does_not_cap_a_team_admin_raising_their_own_team_key(monkeypatch: pytest.MonkeyPatch) -> None: + assert ( + await _denial( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", max_budget=30.0), + existing=_key(user_id=_TEAM_ADMIN, created_by=_TEAM_ADMIN, team_id=_TEAM, max_budget=10.0), + team=_team(["/key/update"]), + caller=_caller(user_id=_TEAM_ADMIN, max_budget=20.0), + ) + is None + ) + + +@pytest.mark.asyncio +async def test_ceiling_caps_ui_session_owner_raise_at_user_budget(monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", max_budget=30.0), + existing=_key(max_budget=10.0), + caller=_caller(team_id=UI_SESSION_TOKEN_TEAM_ID, user_max_budget=20.0), + ) + assert exc.value.status_code == 400 + assert "cannot exceed the caller's own max_budget (20.0)" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_ceiling_respects_personal_key_generation_role_restriction(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.key_generation_settings", + {"personal_key_generation": {"allowed_user_roles": [LitellmUserRoles.PROXY_ADMIN.value]}}, + ) + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", max_budget=None), + existing=_key(max_budget=10.0), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_lower_only_lets_owner_lower_personal_key_budget_without_generate_rights( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "litellm.key_generation_settings", + {"personal_key_generation": {"allowed_user_roles": [LitellmUserRoles.PROXY_ADMIN.value]}}, + ) + assert ( + await _denial( + monkeypatch, + policy="lower_only", + data=UpdateKeyRequest(key="sk-1", max_budget=5.0), + existing=_key(max_budget=10.0), + ) + is None + ) + + +@pytest.mark.asyncio +async def test_ceiling_respects_team_key_generation_role_restriction(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.key_generation_settings", + {"team_key_generation": {"allowed_team_member_roles": ["admin"]}}, + ) + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=_RAISE_WEEKLY_WINDOW, + existing=_team_key_with_weekly(), + team=_team(["/key/update", "/key/generate"]), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["lower_only", "ceiling"]) +async def test_owner_cannot_override_a_cap_inherited_from_a_linked_budget( + monkeypatch: pytest.MonkeyPatch, policy: str +) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy=policy, + data=UpdateKeyRequest(key="sk-1", max_budget=100.0), + existing=_key(max_budget=None, budget_id="shared-tier-10"), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_ceiling_cannot_leave_a_key_uncapped_when_an_upperbound_is_configured( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("litellm.upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(max_budget=20.0)) + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", max_budget=None), + existing=_key(max_budget=10.0), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_ceiling_checks_the_retained_max_budget_when_only_windows_change( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy="ceiling", + data=UpdateKeyRequest(key="sk-1", budget_limits=[]), + existing=_key(max_budget=100.0, budget_limits=[_WEEKLY]), + caller=_caller(max_budget=20.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", [None, "disabled"]) +async def test_owner_cannot_clear_max_budget_by_default(monkeypatch: pytest.MonkeyPatch, policy: str | None) -> None: + with pytest.raises(HTTPException) as exc: + await _update( + monkeypatch, + policy=policy, + data=UpdateKeyRequest(key="sk-1", max_budget=None), + existing=_key(max_budget=10.0), + ) + assert exc.value.status_code == 403 diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index dffd9b9b513..e451c38a251 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -2513,6 +2513,51 @@ async def test_prepare_key_update_data_budget_limits_serializes_windows(): assert windows[0]["reset_at"] is not None +def _key_with_weekly_window(reset_at: datetime) -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token="test-token", + user_id="test-user", + metadata={}, + budget_limits=[{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}], + ) + + +async def _updated_windows_by_duration( + existing_key: LiteLLM_VerificationToken, budget_limits: list[dict[str, object]] +) -> dict[str, dict[str, object]]: + result: Final = await prepare_key_update_data( + data=UpdateKeyRequest.model_validate({"key": "test-token", "budget_limits": budget_limits}), + existing_key_row=existing_key, + ) + return {window["budget_duration"]: window for window in json.loads(result["budget_limits"])} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("weekly_cap", [50.0, 40.0, 60.0]) +async def test_prepare_key_update_data_keeps_a_resubmitted_window_mid_period(weekly_cap: float) -> None: + open_reset_at: Final = datetime.now(timezone.utc) + timedelta(days=3) + + windows: Final = await _updated_windows_by_duration( + _key_with_weekly_window(open_reset_at), + [{"budget_duration": "7d", "max_budget": weekly_cap}, {"budget_duration": "1d", "max_budget": 5.0}], + ) + + assert datetime.fromisoformat(windows["7d"]["reset_at"]) == open_reset_at + assert windows["7d"]["max_budget"] == weekly_cap + assert datetime.fromisoformat(windows["1d"]["reset_at"]) > datetime.now(timezone.utc) + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_restarts_a_resubmitted_window_whose_period_ended() -> None: + stale_reset_at: Final = datetime.now(timezone.utc) - timedelta(days=1) + + windows: Final = await _updated_windows_by_duration( + _key_with_weekly_window(stale_reset_at), [{"budget_duration": "7d", "max_budget": 50.0}] + ) + + assert datetime.fromisoformat(windows["7d"]["reset_at"]) > datetime.now(timezone.utc) + + @pytest.mark.asyncio async def test_prepare_key_update_data_disable_global_guardrails_false_no_premium( monkeypatch, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d697ccac98d..daffc3c0c41 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30213,6 +30213,11 @@ export interface components { * @default false */ search_tool_deny_by_default: boolean; + /** + * Self Serve Budget Policy + * @description Lets a non-admin change max_budget or budget_limits on a key they created and still own via /key/update. 'disabled' (default) requires an admin for any budget change. 'lower_only' lets the owner add or lower limits. 'ceiling' also lets the owner raise or remove limits when they could create that key with the new budget via /key/generate. + */ + self_serve_budget_policy?: ("disabled" | "lower_only" | "ceiling") | null; /** @description Daily check of the spend LiteLLM captured against the provider's own bill (OpenAI via OPENAI_ADMIN_KEY). Publishes litellm_spend_capture_rate per provider and alerts when the ratio over the lookback window falls under the threshold (default 0.9). Off unless set. */ spend_capture_rate_check?: components["schemas"]["SpendCaptureRateCheckSettings"] | null; /** @description Which keys of LiteLLM_SpendLogs.metadata are written to the database. Set exactly one of 'include' (write only these keys) or 'exclude' (drop these keys). 'status' and 'cold_storage_object_key' are always written. Daily spend tables, budgets and logging callbacks still see every key. Unset writes every key */