mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
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 <ryan@berri.ai>
This commit is contained in:
Justin Israel
GitHub
parent
714d84acb6
commit
40e4ade6a6
7 changed files with 874 additions and 79 deletions
|
|
@ -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.",
|
||||
|
|
|
|||
135
litellm/proxy/management_endpoints/key_budget_change.py
Normal file
135
litellm/proxy/management_endpoints/key_budget_change.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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`.
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
467
tests/unit/proxy/management_endpoints/test_key_budget_change.py
Normal file
467
tests/unit/proxy/management_endpoints/test_key_budget_change.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue