diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031120000_team_member_budget_windows/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031120000_team_member_budget_windows/migration.sql new file mode 100644 index 00000000000..89b4b2f8578 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031120000_team_member_budget_windows/migration.sql @@ -0,0 +1,2 @@ +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "team_member_budget_limits" JSONB; +ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN "budget_limits" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 179b741fe7c..e6a46734d74 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -179,6 +179,7 @@ model LiteLLM_ModelTable { // Assign prod keys to groups, not individuals model LiteLLM_TeamTable { + team_member_budget_limits Json? team_id String @id @default(uuid()) team_alias String? organization_id String? @@ -799,6 +800,7 @@ model LiteLLM_UserNotifications { } model LiteLLM_TeamMembership { + budget_limits Json? // Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team user_id String team_id String diff --git a/litellm/models/team.py b/litellm/models/team.py index 8edf10703b1..36edea55594 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -6,11 +6,12 @@ budget-window value types and the team-model alias table). Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ +import enum import json from datetime import datetime -from typing import Final, Literal, Optional +from typing import Annotated, Final, Literal, Optional, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, model_validator from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -50,6 +51,34 @@ class BudgetLimitEntry(LiteLLMPydanticObjectBase): reset_at: datetime | None = None +class TeamMemberBudgetDuration(str, enum.Enum): + HOURLY = "1h" + DAILY = "24h" + WEEKLY = "7d" + MONTHLY = "30d" + + +class TeamMemberBudgetWindow(LiteLLMPydanticObjectBase): + model_config = ConfigDict(frozen=True) + + budget_duration: TeamMemberBudgetDuration + max_budget: float = Field(ge=0, allow_inf_nan=False) + + +def _validate_member_budget_windows( + windows: tuple[TeamMemberBudgetWindow, ...], +) -> tuple[TeamMemberBudgetWindow, ...]: + durations: Final = frozenset(window.budget_duration for window in windows) + if len(durations) != len(windows): + raise ValueError("Each member budget duration must be unique") + return windows + + +TeamMemberBudgetWindows: TypeAlias = Annotated[ + tuple[TeamMemberBudgetWindow, ...], AfterValidator(_validate_member_budget_windows) +] + + class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): id: int | None = None model_aliases: str | dict | None = None @@ -75,6 +104,7 @@ class TeamBase(LiteLLMPydanticObjectBase): max_budget: float | None = None soft_budget: float | None = None budget_duration: str | None = None + team_member_budget_limits: TeamMemberBudgetWindows | None = None budget_limits: list[BudgetLimitEntry] | None = None models: list[str] = [] blocked: bool = False @@ -114,6 +144,7 @@ class LiteLLM_TeamTable(TeamBase): "model_aliases", "router_settings", "budget_limits", + "team_member_budget_limits", ] if isinstance(values, BaseModel): diff --git a/litellm/models/team_membership.py b/litellm/models/team_membership.py index 0ffe8f8dbcc..2b95965b506 100644 --- a/litellm/models/team_membership.py +++ b/litellm/models/team_membership.py @@ -5,18 +5,29 @@ Canonical definition for ``litellm_teammembership``. Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ +import json + +from pydantic import field_validator + from litellm.models.budget import LiteLLM_BudgetTable, LiteLLM_BudgetTableFull +from litellm.models.team import TeamMemberBudgetWindows from litellm.types.llms.base import LiteLLMPydanticObjectBase class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): user_id: str team_id: str + budget_limits: TeamMemberBudgetWindows | None = None budget_id: str | None = None spend: float | None = 0.0 total_spend: float | None = 0.0 litellm_budget_table: LiteLLM_BudgetTableFull | LiteLLM_BudgetTable | None = None + @field_validator("budget_limits", mode="before") + @classmethod + def parse_budget_limits(cls, value: object) -> object: + return json.loads(value) if isinstance(value, str) else value + def safe_get_team_member_rpm_limit(self) -> int | None: if self.litellm_budget_table is not None: return self.litellm_budget_table.rpm_limit diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7b7bfc55ade..2ba42ccd317 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1881,6 +1881,7 @@ class OrgMember(MemberBase): from litellm.models.team import TeamBase as TeamBase # noqa: E402 +from litellm.models.team import TeamMemberBudgetWindows RouterSettingsDict = Annotated[ dict[str, object], @@ -1985,6 +1986,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): object_permission: LiteLLM_ObjectPermissionBase | None = None disable_global_guardrails: bool | None = None require_trace_id: bool | None = None + team_member_budget_limits: TeamMemberBudgetWindows | None = None team_member_budget: float | None = None team_member_budget_duration: str | None = None team_member_model_max_budget: GenericBudgetConfigType | None = None @@ -4259,6 +4261,7 @@ class TeamMemberAddRequest(MemberAddRequest): """ team_id: str = Field(description="The ID of the team to add the member to") + budget_limits: TeamMemberBudgetWindows | None = None max_budget_in_team: float | None = Field( default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", @@ -4278,6 +4281,7 @@ class TeamMemberDeleteRequest(MemberDeleteRequest): class TeamMemberUpdateRequest(TeamMemberDeleteRequest): + budget_limits: TeamMemberBudgetWindows | None = None max_budget_in_team: float | None = None role: Literal["admin", "user"] | None = None tpm_limit: int | None = Field(default=None, description="Tokens per minute limit for this team member") @@ -4317,6 +4321,7 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): class TeamMemberUpdateResponse(MemberUpdateResponse): + budget_limits: TeamMemberBudgetWindows | None = None team_id: str max_budget_in_team: float | None = None tpm_limit: int | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9c0bb93ea16..daf31b42d31 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -7123,6 +7123,35 @@ async def _check_team_member_budget( proxy_logging_obj=proxy_logging_obj, ) loaded_membership = team_membership + from litellm.proxy.common_utils.team_member_budget import MemberBudgetCounter, get_team_member_budget_windows + from litellm.proxy.proxy_server import ensure_window_spend_counter_initialized, get_current_spend + + for window in get_team_member_budget_windows(team_object, loaded_membership): + counter = MemberBudgetCounter.for_member(valid_token.user_id, team_object.team_id, window.budget_duration) + await ensure_window_spend_counter_initialized( + counter_key=counter.key, + entity_type="TeamMember", + entity_id=counter.entity_id, + window_duration=counter.duration.value, + window_start=counter.start, + ) + spend = await get_current_spend( + counter_key=counter.key, + fallback_spend=0.0, + max_budget=window.max_budget, + window_entity_type="TeamMember", + window_entity_id=counter.entity_id, + window_duration=counter.duration.value, + window_start=counter.start, + ) + if spend >= window.max_budget: + raise litellm.BudgetExceededError( + current_cost=spend, + max_budget=window.max_budget, + message=f"User={valid_token.user_id} in Team={team_object.team_id} exceeded {window.budget_duration.value} budget", + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=counter.entity_id, + ) # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. diff --git a/litellm/proxy/common_utils/team_member_budget.py b/litellm/proxy/common_utils/team_member_budget.py new file mode 100644 index 00000000000..42820355ba2 --- /dev/null +++ b/litellm/proxy/common_utils/team_member_budget.py @@ -0,0 +1,61 @@ +"""Concurrent budget windows scoped to one user within one team.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.duration_parser import get_budget_window_start, get_next_standardized_reset_time +from litellm.models.team import LiteLLM_TeamTable, TeamMemberBudgetDuration, TeamMemberBudgetWindow +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_settings + +_MEMBER_ID: Final = TypeAdapter(tuple[str, str]) + + +def member_window_entity_id(user_id: str, team_id: str) -> str: + return _MEMBER_ID.dump_json((user_id, team_id)).decode() + + +def parse_member_window_entity_id(entity_id: str) -> tuple[str, str]: + return _MEMBER_ID.validate_json(entity_id) + + +def get_team_member_budget_windows( + team: LiteLLM_TeamTable, membership: LiteLLM_TeamMembership | None +) -> Sequence[TeamMemberBudgetWindow]: + if membership is not None and membership.budget_limits: + return membership.budget_limits + return team.team_member_budget_limits or () + + +@dataclass(frozen=True, slots=True) +class MemberBudgetCounter: + entity_id: str + duration: TeamMemberBudgetDuration + start: datetime + reset_at: datetime + + @property + def key(self) -> str: + return f"spend:team_member_window:{self.entity_id}:{self.duration.value}:{int(self.start.timestamp())}" + + @classmethod + def for_member( + cls, user_id: str, team_id: str, duration: TeamMemberBudgetDuration, now: datetime | None = None + ) -> "MemberBudgetCounter": + settings: Final = get_budget_reset_settings() + reset_at: Final = get_next_standardized_reset_time( + duration=duration.value, + current_time=now if now is not None else datetime.now(timezone.utc), + timezone_str=settings.timezone, + reset_time_of_day=settings.reset_time_of_day, + ) + return cls( + entity_id=member_window_entity_id(user_id, team_id), + duration=duration, + start=get_budget_window_start(duration.value, reset_at), + reset_at=reset_at, + ) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 99b49331932..7a6dac07860 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -31,7 +31,10 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( ) if TYPE_CHECKING: + from prisma import models as prisma_models + from litellm.proxy.utils import PrismaClient + from litellm.repositories.prisma_protocols import TableActions _SELECT_EXISTING_ROWS_SQL: Final = ( @@ -133,6 +136,9 @@ class _UpsertBatchManager(Protocol): class _UpsertTransaction(Protocol): + @property + def litellm_budgetwindowspend(self) -> "TableActions[prisma_models.LiteLLM_BudgetWindowSpend]": ... + def batch_(self) -> _UpsertBatchManager: ... @@ -204,6 +210,26 @@ async def spend_logs_seed_totals( counts everything: that can only over-count once, which enforcement tolerates, whereas under-counting is a budget bypass. """ + if entity_type == Litellm_EntityType.TEAM_MEMBER.value: + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + total: Final = await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, entity_type="TeamMember", entity_id=entity_id, window_start=window_start + ) + before: Final = ( + total + if batch_started_at is None + else await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, + entity_type="TeamMember", + entity_id=entity_id, + window_start=window_start, + window_end=_exclusion_upper_bound(batch_started_at), + ) + ) + if total is None or before is None: + return None + return WindowSeedTotals(total=total, before_batch=before) if entity_type == Litellm_EntityType.KEY.value: bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL elif entity_type == Litellm_EntityType.TEAM.value: @@ -342,10 +368,31 @@ async def commit_window_spend_updates( if not transactions: return - ordered: Final = tuple(sorted(transactions, key=window_spend_group_key)) - existing_primary_keys: Final = await _existing_primary_keys( - prisma_client=prisma_client, - transactions=ordered, + member_transactions: Final = tuple( + sorted( + (t for t in transactions if t["entity_type"] == Litellm_EntityType.TEAM_MEMBER.value), + key=window_spend_group_key, + ) + ) + member_seed_bases: Final = tuple( + [ + await _member_window_seed_base(prisma_client, transaction, spend_logs_aggregate) + for transaction in member_transactions + ] + ) + ordered: Final = tuple( + sorted( + (t for t in transactions if t["entity_type"] != Litellm_EntityType.TEAM_MEMBER.value), + key=window_spend_group_key, + ) + ) + existing_primary_keys: Final[frozenset[tuple[str, str, str]]] = ( + await _existing_primary_keys( + prisma_client=prisma_client, + transactions=ordered, + ) + if ordered + else frozenset() ) seed_bases: Final = tuple( [ @@ -368,13 +415,71 @@ async def commit_window_spend_updates( async with ( db_span("commit_window_spend_updates", "LiteLLM_BudgetWindowSpend"), _upsert_transaction(prisma_client) as db_transaction, - db_transaction.batch_() as batcher, ): - for transaction, seed_base in zip(ordered, seed_bases): - batcher.execute_raw( - _UPSERT_WINDOW_SPEND_SQL, - *_upsert_params(transaction=transaction, seed_base=seed_base, now=now), - ) + for transaction, seed_base in zip(member_transactions, member_seed_bases): + await _write_member_window_spend(db_transaction.litellm_budgetwindowspend, transaction, seed_base) + async with db_transaction.batch_() as batcher: + for transaction, seed_base in zip(ordered, seed_bases): + batcher.execute_raw( + _UPSERT_WINDOW_SPEND_SQL, + *_upsert_params(transaction=transaction, seed_base=seed_base, now=now), + ) + + +async def _member_window_seed_base( + prisma_client: "PrismaClient", + transaction: WindowSpendTransaction, + spend_logs_aggregate: WindowSpendLogsAggregate, +) -> float: + from litellm.repositories.table_repositories import BudgetWindowSpendRepository + + existing: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( + where={ + "entity_type_entity_id_window_duration": { + "entity_type": transaction["entity_type"], + "entity_id": transaction["entity_id"], + "window_duration": transaction["window_duration"], + } + } + ) + if existing is not None: + return 0.0 + totals: Final = await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), + batch_started_at=_transaction_started_at(transaction), + ) + if totals is None: + return 0.0 + return _seed_base(totals, transaction["spend"]) + + +async def _write_member_window_spend( + table: "TableActions[prisma_models.LiteLLM_BudgetWindowSpend]", + transaction: WindowSpendTransaction, + seed: float, +) -> None: + key: Final = { + "entity_type": transaction["entity_type"], + "entity_id": transaction["entity_id"], + "window_duration": transaction["window_duration"], + } + start: Final = datetime.fromisoformat(transaction["window_start"]) + # Keep this row locked until the entire flush commits, including key/team windows. + await table.upsert( + where={"entity_type_entity_id_window_duration": key}, + data={"create": {**key, "window_start": start, "spend": seed}, "update": {"spend": {"increment": 0.0}}}, + ) + await table.update_many( + where={**key, "window_start": {"lt": start}}, + data={"window_start": start, "spend": seed}, + ) + await table.update( + where={"entity_type_entity_id_window_duration": key}, + data={"spend": {"increment": transaction["spend"]}}, + ) async def roll_window_spend_row( diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..95648e18081 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -55,6 +55,7 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( { "Key": Litellm_EntityType.KEY.value, "Team": Litellm_EntityType.TEAM.value, + "TeamMember": Litellm_EntityType.TEAM_MEMBER.value, } ) @@ -68,6 +69,18 @@ _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( ) +def _window_spend_log_identity(entity_type: str, entity_id: str) -> Mapping[str, str] | None: + if entity_type == "TeamMember": + from litellm.proxy.common_utils.team_member_budget import parse_member_window_entity_id + + user_id, team_id = parse_member_window_entity_id(entity_id) + return {"user": user_id, "team_id": team_id} + group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) + if group_field is None: + return None + return {group_field: entity_id} + + def _as_utc(value: datetime) -> datetime: return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) @@ -388,21 +401,20 @@ class SpendCounterReseed: entity_type: str, entity_id: str, window_start: datetime, + window_end: datetime | None = None, ) -> float | None: if prisma_client is None: return None - group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) - if group_field is None: + identity: Final = _window_spend_log_identity(entity_type, entity_id) + if identity is None: return None - where: Final = { - group_field: entity_id, - "startTime": {"gte": window_start}, - } - + group_fields: Final = tuple(identity) + period: Final = {"gte": window_start, **({"lt": window_end} if window_end is not None else {})} + where: Final = {**identity, "startTime": period} try: response: Final = await SpendLogsRepository(prisma_client).table.group_by( - by=[group_field], + by=group_fields, where=where, sum={"spend": True}, ) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 983c3d24513..59ff98614a3 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,6 +1,6 @@ import json import math -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Optional, Union @@ -44,6 +44,8 @@ def validate_budget_duration(budget_duration: str | None, status_code: int = 400 from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache +from litellm.models.team import TeamMemberBudgetWindow +from litellm.models.team_membership import LiteLLM_TeamMembership from litellm.proxy._types import ( # re-exported CommonProxyErrors, KeyRequestBase, @@ -69,8 +71,11 @@ from litellm.repositories.team_repository import TeamRepository from litellm.types.utils import BudgetConfig if TYPE_CHECKING: + from prisma import models as prisma_models + from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest from litellm.proxy.utils import PrismaClient, ProxyLogging + from litellm.repositories.prisma_protocols import TableActions _user_has_admin_view: Final = user_api_key_has_admin_view @@ -625,6 +630,24 @@ def _has_meaningful_budget_limit(budget_values: Mapping[str, object]) -> bool: return any(_is_set_budget_value(budget_values.get(field)) for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS) +async def upsert_team_member_budget_windows( + table: "TableActions[prisma_models.LiteLLM_TeamMembership]", + team_id: str, + user_id: str, + windows: Sequence[TeamMemberBudgetWindow], +) -> LiteLLM_TeamMembership: + payload: Final = "[" + ",".join(window.model_dump_json() for window in windows) + "]" + membership: Final = await table.upsert( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + data={ + "create": {"user_id": user_id, "team_id": team_id, "budget_limits": payload}, + "update": {"budget_limits": payload}, + }, + include={"litellm_budget_table": True}, + ) + return LiteLLM_TeamMembership.model_validate(membership.model_dump()) + + async def upsert_budget_and_membership( tx, *, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 61da5599886..3f2328c5dae 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -8074,7 +8074,7 @@ async def block_key( --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ --header 'Content-Type: application/json' \ --data '{ - "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + "key": "" }' ``` @@ -8188,7 +8188,7 @@ async def unblock_key( --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ --header 'Content-Type: application/json' \ --data '{ - "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + "key": "" }' ``` diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 0ece14eb90a..fd072086599 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -48,6 +48,7 @@ SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset( "metadata", "team_member_budget", "team_member_budget_duration", + "team_member_budget_limits", "team_member_tpm_limit", "team_member_rpm_limit", } diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b712ba42a6a..34111dea0f3 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -171,6 +171,7 @@ from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # l team_member_has_permission, update_metadata_fields, upsert_budget_and_membership, + upsert_team_member_budget_windows, user_api_key_has_admin_view, validate_budget_duration, validate_team_model_max_budget, @@ -533,7 +534,14 @@ def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]" ) -def _team_membership_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_TeamMembership]": +def _team_membership_db( + prisma_client: PrismaClient | None, + tx: "MemberWriteTx | Prisma | None" = None, +) -> "TableActions[prisma_models.LiteLLM_TeamMembership]": + if tx is not None: + return cast( # cast-ok: Prisma JSON writes are strings; TableActions widens inputs at the transaction boundary + "TableActions[prisma_models.LiteLLM_TeamMembership]", tx.litellm_teammembership + ) return TeamMembershipRepository(prisma_client).table @@ -1961,10 +1969,7 @@ async def new_team( ) ## ADD TO TEAM TABLE - complete_team_data: Final = LiteLLM_TeamTable( - **data_json, - model_id=_model_id, - ) + complete_team_data: Final = LiteLLM_TeamTable.model_validate({**data_json, "model_id": _model_id}) # Set Management Endpoint Metadata Fields for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: @@ -3147,6 +3152,8 @@ async def _validate_team_member_add_permissions( # budget and model controls stay admin-only. Reject them here so a # self-joining non-admin cannot set their own cap, reset window, or model # scope via the bypass. + if data.budget_limits is not None: + raise HTTPException(status_code=403, detail="Only admins may set member budget windows") if data.max_budget_in_team is not None or data.budget_duration is not None or data.allowed_models is not None: raise HTTPException( status_code=403, @@ -3238,6 +3245,13 @@ async def _process_team_members( status_code=500, detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e}"}, ) + if data.budget_limits is not None: + updated_tm = await upsert_team_member_budget_windows( + table=_team_membership_db(prisma_client, tx), + team_id=data.team_id, + user_id=updated_user.user_id, + windows=data.budget_limits, + ) updated_users.append(updated_user) if updated_tm is not None: updated_team_memberships.append(updated_tm) @@ -4233,7 +4247,14 @@ async def team_member_update( budget_patch=budget_patch, team_default_budget_id=team_default_budget_id, ) - if budget_patch: + if data.budget_limits is not None: + await upsert_team_member_budget_windows( + table=_team_membership_db(prisma_client, tx), + team_id=data.team_id, + user_id=received_user_id, + windows=data.budget_limits, + ) + if budget_patch or data.budget_limits is not None: await invalidate_team_member_spend_state( user_id=received_user_id, team_id=data.team_id, @@ -4257,6 +4278,7 @@ async def team_member_update( user_id=received_user_id, user_email=data.user_email, max_budget_in_team=data.max_budget_in_team, + budget_limits=data.budget_limits, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, budget_duration=data.budget_duration, @@ -5185,6 +5207,7 @@ def _transform_teams_to_deleted_records( "litellm_model_table", "object_permission", "id", + "team_member_budget_limits", "budget_limits", # not in LiteLLM_DeletedTeamTable schema "default_team_member_models", # not in LiteLLM_DeletedTeamTable schema ): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 634c709b3ca..2e53d76cd0f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3753,15 +3753,63 @@ async def _increment_spend_counters_batched( scope_user_id: str, scope_team_id: str ) -> tuple[PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" - if team_member_counter_key in reserved_counter_keys: - return () - return ( - await _prepare_spend_counter_increment( - counter_key=team_member_counter_key, - source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", - increment=cost, - ), + primary_pending: Final = ( + () + if team_member_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=team_member_counter_key, + source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", + increment=cost, + ), + ) ) + if prisma_client is None: + return primary_pending + from litellm.models.team import TeamMemberBudgetWindow + from litellm.proxy.auth.auth_checks import get_team_membership, get_team_object + from litellm.proxy.common_utils.team_member_budget import MemberBudgetCounter, get_team_member_budget_windows + + team: Final = await get_team_object( + team_id=scope_team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache + ) + membership: Final = await get_team_membership( + user_id=scope_user_id, + team_id=scope_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + async def member_window_increment(window: TeamMemberBudgetWindow) -> PendingSpendIncrement | None: + counter: Final = MemberBudgetCounter.for_member(scope_user_id, scope_team_id, window.budget_duration) + pending: Final = ( + await _prepare_window_spend_counter_increment( + counter_key=counter.key, + entity_type="TeamMember", + entity_id=counter.entity_id, + window_duration=counter.duration.value, + window_start=counter.start, + increment=cost, + ) + if counter.key not in reserved_counter_keys + else None + ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.TEAM_MEMBER, + entity_id=counter.entity_id, + reset_at=counter.reset_at, + window_duration=counter.duration.value, + window_start=counter.start, + increment=cost, + request_started_at=request_started_at, + ) + return pending + + window_pending: Final = await asyncio.gather( + *(member_window_increment(window) for window in get_team_member_budget_windows(team, membership)), + return_exceptions=True, + ) + return primary_pending + tuple(item for item in window_pending if item is not None) async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 179b741fe7c..e6a46734d74 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -179,6 +179,7 @@ model LiteLLM_ModelTable { // Assign prod keys to groups, not individuals model LiteLLM_TeamTable { + team_member_budget_limits Json? team_id String @id @default(uuid()) team_alias String? organization_id String? @@ -799,6 +800,7 @@ model LiteLLM_UserNotifications { } model LiteLLM_TeamMembership { + budget_limits Json? // Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team user_id String team_id String diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index fc6fb266d04..01fc87c08d1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -582,6 +582,24 @@ async def _get_budget_counters( ) if team_member_counter is not None: counters.append(team_member_counter) + if team_object is not None and valid_token.user_id is not None: + from litellm.proxy.common_utils.team_member_budget import MemberBudgetCounter, get_team_member_budget_windows + + membership: Final = await _cached_team_membership(valid_token.user_id, team_object.team_id, user_api_key_cache) + for window in get_team_member_budget_windows(team_object, membership): + counter = MemberBudgetCounter.for_member(valid_token.user_id, team_object.team_id, window.budget_duration) + counters.append( + _BudgetCounter( + counter_key=counter.key, + max_budget=window.max_budget, + fallback_spend=0.0, + entity_type="TeamMember", + entity_id=counter.entity_id, + spend_log_entity_id=counter.entity_id, + window_duration=counter.duration.value, + window_start=counter.start, + ) + ) org_counter: Final = await _get_org_budget_counter( valid_token=valid_token, @@ -726,6 +744,21 @@ def _dedupe_tags(tags: list[str]) -> list[str]: @with_service_target(AUTH_OBJECTS_TARGET) +async def _cached_team_membership( + user_id: str, team_id: str, user_api_key_cache: UserApiKeyCache +) -> LiteLLM_TeamMembership | None: + cached: Final = _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) + ) + ) + if isinstance(cached, LiteLLM_TeamMembership): + return cached + if isinstance(cached, dict): + return LiteLLM_TeamMembership.model_validate(cached) + return None + + async def _get_team_member_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, @@ -738,14 +771,7 @@ async def _get_team_member_budget_counter( membership_cache_key: Final = team_membership_reservation_cache_key( user_id=valid_token.user_id, team_id=team_object.team_id ) - cached_team_membership: Final = _CACHED_VALUE.validate_python( - await user_api_key_cache.async_get_cache(key=membership_cache_key) - ) - team_membership: LiteLLM_TeamMembership | None = None - if isinstance(cached_team_membership, LiteLLM_TeamMembership): - team_membership = cached_team_membership - elif isinstance(cached_team_membership, dict): - team_membership = LiteLLM_TeamMembership(**cached_team_membership) + team_membership: Final = await _cached_team_membership(valid_token.user_id, team_object.team_id, user_api_key_cache) member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None now: Final = datetime.now(timezone.utc) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a1b36e5b7a1..7339cd56feb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6331,12 +6331,11 @@ class PrismaClient: raise e def jsonify_team_object(self, db_data: Mapping[str, object]) -> dict[str, object]: - db_data = self.jsonify_object(data=db_data) - if db_data.get("members_with_roles", None) is not None and isinstance(db_data["members_with_roles"], list): - db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"]) - if db_data.get("budget_limits", None) is not None and isinstance(db_data["budget_limits"], list): - db_data["budget_limits"] = json.dumps(db_data["budget_limits"]) - return db_data + serialized: Final = self.jsonify_object(data=db_data) + for column in ("members_with_roles", "budget_limits", "team_member_budget_limits"): + if isinstance(serialized.get(column), (list, tuple)): + serialized[column] = json.dumps(serialized[column]) + return serialized # Define a retrying strategy with exponential backoff @backoff.on_exception( diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index b89d87aa354..9aa26ef9a23 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -82,6 +82,7 @@ _JSON_ENCODED_TEAM_FIELDS: Final = ( "model_max_budget", "router_settings", "budget_limits", + "team_member_budget_limits", "members_with_roles", ) diff --git a/schema.prisma b/schema.prisma index 179b741fe7c..e6a46734d74 100644 --- a/schema.prisma +++ b/schema.prisma @@ -179,6 +179,7 @@ model LiteLLM_ModelTable { // Assign prod keys to groups, not individuals model LiteLLM_TeamTable { + team_member_budget_limits Json? team_id String @id @default(uuid()) team_alias String? organization_id String? @@ -799,6 +800,7 @@ model LiteLLM_UserNotifications { } model LiteLLM_TeamMembership { + budget_limits Json? // Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team user_id String team_id String diff --git a/tests/integration/database/test_member_window_spend_atomicity.py b/tests/integration/database/test_member_window_spend_atomicity.py new file mode 100644 index 00000000000..b327306cb87 --- /dev/null +++ b/tests/integration/database/test_member_window_spend_atomicity.py @@ -0,0 +1,59 @@ +"""A retried flush must not charge a member again after a sibling write fails.""" + +import json +import os +import uuid +from datetime import datetime, timezone + +import psycopg +import pytest +from psycopg import sql +from prisma.errors import RawQueryError + +from integration._support.prisma_clients import connected_proxy_prisma +from litellm.proxy.db.budget_window_spend_writer import commit_window_spend_updates +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import build_window_spend_transaction + + +@pytest.mark.asyncio +async def test_failed_key_window_rolls_back_member_spend_before_retry() -> None: + marker = uuid.uuid4().hex + member_id = json.dumps(["member-" + marker, "team-" + marker], separators=(",", ":")) + key_id = "key-" + marker + constraint = "window_fault_" + marker + start = datetime(2026, 10, 1, tzinfo=timezone.utc) + transactions = [ + build_window_spend_transaction("team_member", member_id, "30d", start, 0.06, start), + build_window_spend_transaction("key", key_id, "30d", start, 0.06, start), + ] + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as control: + control.execute( + sql.SQL('ALTER TABLE "LiteLLM_BudgetWindowSpend" ADD CONSTRAINT {} CHECK (entity_id <> {})').format( + sql.Identifier(constraint), + sql.Literal(key_id), + ) + ) + try: + async with connected_proxy_prisma() as client: + with pytest.raises(RawQueryError, match=constraint): + await commit_window_spend_updates(client, transactions) + assert ( + await client.db.litellm_budgetwindowspend.find_many( + where={"entity_id": {"in": [member_id, key_id]}} + ) + == [] + ) + finally: + control.execute( + sql.SQL('ALTER TABLE "LiteLLM_BudgetWindowSpend" DROP CONSTRAINT {}').format(sql.Identifier(constraint)) + ) + try: + async with connected_proxy_prisma() as client: + await commit_window_spend_updates(client, transactions) + rows = await client.db.litellm_budgetwindowspend.find_many( + where={"entity_id": {"in": [member_id, key_id]}} + ) + assert len(rows) == 2 + assert all(row.spend == pytest.approx(0.06) for row in rows) + finally: + control.execute('DELETE FROM "LiteLLM_BudgetWindowSpend" WHERE entity_id = ANY(%s)', ([member_id, key_id],)) diff --git a/tests/integration/spend/test_team_member_budget_windows.py b/tests/integration/spend/test_team_member_budget_windows.py new file mode 100644 index 00000000000..1ecd4d95643 --- /dev/null +++ b/tests/integration/spend/test_team_member_budget_windows.py @@ -0,0 +1,162 @@ +"""Real proxy, database and Redis: concurrent budgets belong to a member/team pair.""" + +import uuid +from typing import Final + +import pytest + +from integration._support.client import Gateway, eventually + +pytestmark: Final = pytest.mark.shares_worker_state +CALL_COST: Final = 20 * 0.001 + 20 * 0.002 + + +def test_member_windows_persist_enforce_and_restore_inheritance(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + defaults = [{"budget_duration": "24h", "max_budget": 0.05}, {"budget_duration": "30d", "max_budget": 0.5}] + team = scenario.team(models=[model], team_member_budget_limits=defaults) + alice = scenario.user(auto_create_key=False) + bob = scenario.user(auto_create_key=False) + for user in (alice, bob): + gateway.post( + "/team/member_add", + { + "team_id": team, + "member": {"user_id": user, "role": "user"}, + **({"max_budget_in_team": 10, "budget_duration": "30d"} if user == alice else {}), + }, + ) + alice_key = scenario.key(team_id=team, user_id=alice) + bob_key = scenario.key(team_id=team, user_id=bob) + saved = gateway.get("/team/info", {"team_id": team}) + assert saved["team_info"]["team_member_budget_limits"] == defaults + + def chat(key: str): + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": uuid.uuid4().hex}], + "max_tokens": 1, + }, + key=key, + ) + + first = chat(alice_key) + assert first.status_code == 200, first.text + assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(CALL_COST) + eventually( + lambda: gateway.rows( + 'SELECT spend FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type = %s AND entity_id = %s', + ("team_member", f'["{alice}","{team}"]'), + ), + lambda rows: len(rows) == 2 and all(row["spend"] == pytest.approx(CALL_COST) for row in rows), + seconds=70, + ) + capped = chat(alice_key) + assert capped.status_code == 422 and "24h" in capped.text, capped.text + assert chat(bob_key).status_code == 200 + + other_team = scenario.team(models=[model], team_member_budget_limits=defaults) + gateway.post("/team/member_add", {"team_id": other_team, "member": {"user_id": alice, "role": "user"}}) + other_key = scenario.key(team_id=other_team, user_id=alice) + assert chat(other_key).status_code == 200 + + custom = [{"budget_duration": "24h", "max_budget": 1}, {"budget_duration": "30d", "max_budget": 0.05}] + gateway.post("/team/member_update", {"team_id": team, "user_id": alice, "budget_limits": custom}) + saved_member = gateway.get("/team/info", {"team_id": team}) + alice_membership = next(m for m in saved_member["team_memberships"] if m["user_id"] == alice) + assert alice_membership["budget_limits"] == custom + assert alice_membership["litellm_budget_table"]["max_budget"] == 10 + assert alice_membership["litellm_budget_table"]["budget_duration"] == "30d" + monthly = chat(alice_key) + assert monthly.status_code == 422 and "30d" in monthly.text, monthly.text + + gateway.post("/team/member_update", {"team_id": team, "user_id": alice, "budget_limits": []}) + inherited = chat(alice_key) + assert inherited.status_code == 422 and "24h" in inherited.text, inherited.text + + gateway.post("/team/update", {"team_id": team, "team_member_budget_limits": []}) + assert chat(alice_key).status_code == 200 + + +@pytest.mark.asyncio +async def test_expired_daily_window_resets_without_resetting_monthly_spend(gateway: Gateway) -> None: + import json + from datetime import datetime, timedelta, timezone + + from integration._support.prisma_clients import connected_prisma + + with gateway.scenario() as scenario: + model = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + defaults = [{"budget_duration": "24h", "max_budget": 0.05}, {"budget_duration": "30d", "max_budget": 0.1}] + team = scenario.team(models=[model], team_member_budget_limits=defaults) + member = scenario.user(auto_create_key=False) + gateway.post("/team/member_add", {"team_id": team, "member": {"user_id": member, "role": "user"}}) + key = scenario.key(team_id=team, user_id=member) + today = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + identity = json.dumps([member, team], separators=(",", ":")) + async with connected_prisma() as db: + await db.litellm_budgetwindowspend.create( + data={ + "entity_type": "team_member", + "entity_id": identity, + "window_duration": "24h", + "window_start": today - timedelta(days=1), + "spend": 10.0, + } + ) + await db.litellm_budgetwindowspend.create( + data={ + "entity_type": "team_member", + "entity_id": identity, + "window_duration": "30d", + "window_start": today.replace(day=1), + "spend": 0.2, + } + ) + + def chat(): + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": uuid.uuid4().hex}], + "max_tokens": 1, + }, + key=key, + ) + + monthly = chat() + assert monthly.status_code == 422 and "30d" in monthly.text, monthly.text + gateway.post( + "/team/update", + { + "team_id": team, + "team_member_budget_limits": [ + defaults[0], + {"budget_duration": "30d", "max_budget": 1}, + ], + }, + ) + renewed = chat() + assert renewed.status_code == 200, renewed.text + eventually( + lambda: gateway.rows( + 'SELECT window_duration, spend FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type = %s AND entity_id = %s', + ("team_member", identity), + ), + lambda rows: ( + {row["window_duration"]: row["spend"] for row in rows} + == { + "24h": pytest.approx(CALL_COST), + "30d": pytest.approx(0.2 + CALL_COST), + } + ), + seconds=70, + ) + daily = chat() + assert daily.status_code == 422 and "24h" in daily.text, daily.text diff --git a/tests/unit/proxy/auth/test_team_member_budget.py b/tests/unit/proxy/auth/test_team_member_budget.py index 6055e4f0356..fa155b19b3f 100644 --- a/tests/unit/proxy/auth/test_team_member_budget.py +++ b/tests/unit/proxy/auth/test_team_member_budget.py @@ -469,21 +469,56 @@ def test_resolve_effective_team_member_budget_preserves_member_and_default_prece assert resolve_effective_team_member_budget(member_budget, None, now) == 110 assert resolve_effective_team_member_budget(expired_member_budget, None, now) == 100 - assert resolve_effective_team_member_budget( - active_inherited_increase, - LiteLLM_BudgetTable(max_budget=50), - now, - ) == 60 - assert resolve_effective_team_member_budget( - expired_inherited_increase, - LiteLLM_BudgetTable(max_budget=50), - now, - ) == 50 - assert resolve_effective_team_member_budget( - LiteLLM_BudgetTable(max_budget=0), - LiteLLM_BudgetTable(max_budget=50), - now, - ) == 0 + assert ( + resolve_effective_team_member_budget( + active_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) + == 60 + ) + assert ( + resolve_effective_team_member_budget( + expired_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) + == 50 + ) + assert ( + resolve_effective_team_member_budget( + LiteLLM_BudgetTable(max_budget=0), + LiteLLM_BudgetTable(max_budget=50), + now, + ) + == 0 + ) assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=50), now) == 50 assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=0), now) is None assert resolve_effective_team_member_budget(None, None, now) is None + + +@pytest.mark.asyncio +async def test_jwt_member_without_user_row_still_enforces_default_windows() -> None: + from litellm.models.team import TeamMemberBudgetDuration, TeamMemberBudgetWindow + from litellm.proxy.auth.auth_checks import _check_team_member_budget + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + team = LiteLLM_TeamTable( + team_id="jwt-member-budget-team", + team_member_budget_limits=( + TeamMemberBudgetWindow(budget_duration=TeamMemberBudgetDuration.DAILY, max_budget=0), + ), + ) + with pytest.raises(litellm.BudgetExceededError, match="24h"): + await _check_team_member_budget( + team_object=team, + user_object=None, + valid_token=UserAPIKeyAuth(user_id="jwt-member-only", team_id=team.team_id), + prisma_client=None, + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + team_membership_loaded=True, + ) diff --git a/tests/unit/proxy/common_utils/test_team_member_budget.py b/tests/unit/proxy/common_utils/test_team_member_budget.py new file mode 100644 index 00000000000..7a229b0011d --- /dev/null +++ b/tests/unit/proxy/common_utils/test_team_member_budget.py @@ -0,0 +1,85 @@ +from datetime import datetime, timezone + +import pytest +from pydantic import ValidationError + +from litellm.models.team import LiteLLM_TeamTable, TeamMemberBudgetDuration, TeamMemberBudgetWindow +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.proxy._types import NewTeamRequest, TeamMemberAddRequest, TeamMemberUpdateRequest, UpdateTeamRequest +from litellm.proxy.common_utils.team_member_budget import ( + MemberBudgetCounter, + get_team_member_budget_windows, + member_window_entity_id, + parse_member_window_entity_id, +) + + +@pytest.mark.parametrize("amount", [-1, float("inf"), float("nan")]) +def test_rejects_invalid_caps(amount: float) -> None: + with pytest.raises(ValidationError): + TeamMemberBudgetWindow(budget_duration=TeamMemberBudgetDuration.DAILY, max_budget=amount) + + +@pytest.mark.parametrize( + "request_type, fields, field", + [ + (NewTeamRequest, {}, "team_member_budget_limits"), + (UpdateTeamRequest, {"team_id": "team"}, "team_member_budget_limits"), + (TeamMemberAddRequest, {"team_id": "team", "member": {"user_id": "member", "role": "user"}}, "budget_limits"), + (TeamMemberUpdateRequest, {"team_id": "team", "user_id": "member"}, "budget_limits"), + ], +) +def test_rejects_duplicate_windows(request_type: type, fields: dict, field: str) -> None: + duplicate = [{"budget_duration": "24h", "max_budget": 20}] * 2 + with pytest.raises(ValidationError, match="unique"): + request_type(**fields, **{field: duplicate}) + + +def test_override_and_clear_preserve_team_inheritance() -> None: + defaults = [TeamMemberBudgetWindow(budget_duration=TeamMemberBudgetDuration.DAILY, max_budget=20)] + custom = [TeamMemberBudgetWindow(budget_duration=TeamMemberBudgetDuration.DAILY, max_budget=0)] + team = LiteLLM_TeamTable(team_id="team", team_member_budget_limits=defaults) + member = LiteLLM_TeamMembership(user_id="member", team_id="team", budget_limits=custom) + assert get_team_member_budget_windows(team, member) == tuple(custom) + member.budget_limits = [] + assert get_team_member_budget_windows(team, member) == tuple(defaults) + assert get_team_member_budget_windows(team, None) == tuple(defaults) + assert get_team_member_budget_windows(LiteLLM_TeamTable(team_id="off"), None) == () + + +def test_reset_boundaries_are_independent() -> None: + before = datetime(2026, 10, 10, 23, 59, tzinfo=timezone.utc) + after = datetime(2026, 10, 11, 0, 1, tzinfo=timezone.utc) + daily = MemberBudgetCounter.for_member("member", "team", TeamMemberBudgetDuration.DAILY, before) + monthly = MemberBudgetCounter.for_member("member", "team", TeamMemberBudgetDuration.MONTHLY, before) + assert MemberBudgetCounter.for_member("member", "team", daily.duration, after).key != daily.key + assert MemberBudgetCounter.for_member("member", "team", monthly.duration, after).key == monthly.key + assert MemberBudgetCounter.for_member("other", "team", daily.duration, before).key != daily.key + assert MemberBudgetCounter.for_member("member", "other", daily.duration, before).key != daily.key + + +def test_member_identity_cannot_collide_at_a_separator() -> None: + first = member_window_entity_id("user:other", "team") + second = member_window_entity_id("user", "other:team") + assert first != second + assert parse_member_window_entity_id(first) == ("user:other", "team") + + +@pytest.mark.parametrize("request_type", [NewTeamRequest, UpdateTeamRequest]) +def test_default_windows_round_trip_through_prisma_json(request_type: type) -> None: + import json + from litellm.proxy.utils import PrismaClient + + data = request_type( + team_id="team", + team_member_budget_limits=[ + {"budget_duration": "24h", "max_budget": 20}, + {"budget_duration": "30d", "max_budget": 200}, + ], + ) + client = object.__new__(PrismaClient) + serialized = client.jsonify_team_object(data.model_dump(exclude_none=True)) + windows = json.loads(serialized["team_member_budget_limits"]) + assert windows == [{"budget_duration": "24h", "max_budget": 20}, {"budget_duration": "30d", "max_budget": 200}] + restored = LiteLLM_TeamTable(team_id="team", team_member_budget_limits=serialized["team_member_budget_limits"]) + assert restored.team_member_budget_limits == data.team_member_budget_limits diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx index 34fc88e04bb..2fab2059393 100644 --- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx @@ -1,7 +1,7 @@ import { describe, expect, it, vi } from "vitest"; import userEvent from "@testing-library/user-event"; -import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; +import { act, chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; import TeamAdminSettingsForm, { type TeamAdminMetadataSettings } from "./TeamAdminSettingsForm"; @@ -301,3 +301,24 @@ describe("TeamAdminSettingsForm", () => { }); }); }); + +describe("team admin budget windows form", () => { + it("should configure windows with the same optional editor when permitted", async () => { + const user = userEvent.setup(); + const { onSave } = renderForm(new Set(["team_member_budget_limits"])); + const toggle = screen.getByRole("switch", { name: "Enable budget windows" }); + expect(toggle).not.toBeChecked(); + await user.click(toggle); + act(() => + fireEvent.change(screen.getByRole("spinbutton", { name: "Window 1 budget (USD)" }), { target: { value: "20" } }), + ); + await user.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => + expect(onSave).toHaveBeenCalledWith({ team_member_budget_limits: [{ budget_duration: "24h", max_budget: 20 }] }), + ); + }); + it("should hide windows when the proxy has not granted permission", () => { + renderForm(new Set(["tpm_limit"])); + expect(screen.queryByRole("switch", { name: "Enable budget windows" })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx index 707d4181a4a..58be0998804 100644 --- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx @@ -19,6 +19,8 @@ import MetadataKeyValueFields, { metadataPairsSchema, } from "../common_components/MetadataKeyValueFields"; import NumericalInput from "../shared/numerical_input"; +import { MemberBudgetWindows } from "./MemberBudgetWindows"; +import { isMemberBudgetWindows, memberBudgetWindowsSchema } from "./memberBudgetWindowSchema"; import { isTeamAdminDurationSettingsField, TEAM_ADMIN_METADATA_FIELD, @@ -42,6 +44,7 @@ const teamAdminSettingsShape = { soft_budget: numericInputSchema, team_member_budget: numericInputSchema, team_member_budget_duration: z.string().nullish(), + team_member_budget_limits: memberBudgetWindowsSchema.optional(), team_member_tpm_limit: numericInputSchema, team_member_rpm_limit: numericInputSchema, }; @@ -128,7 +131,16 @@ function TeamAdminSettingsFormFields({ const form = useZodForm(schema, { defaultValues: { ...initialValues, metadata: [...(metadata?.pairs ?? [])] } }); const draft = useWatch({ control: form.control }); const changesFrom = (values: typeof draft): TeamAdminSettingsChanges => ({ - ...teamAdminSettingsChanges(values, initialValues, editableFields), + ...teamAdminSettingsChanges( + { + ...values, + team_member_budget_limits: isMemberBudgetWindows(values.team_member_budget_limits) + ? values.team_member_budget_limits + : undefined, + }, + initialValues, + editableFields, + ), ...(editsMetadata ? teamAdminMetadataChange(values.metadata, metadata.context) : {}), }); const hasChanges = Object.keys(changesFrom(draft)).length > 0; @@ -170,6 +182,15 @@ function TeamAdminSettingsFormFields({ } ))} + {editableFields.has("team_member_budget_limits") && ( + + {({ value, onChange }) => } + + )} {editsMetadata && ( {teamAdminFieldLabel(TEAM_ADMIN_METADATA_FIELD)} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index e22a7aee381..197943c3005 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1334,6 +1334,7 @@ const TeamInfoView: React.FC = ({ soft_budget: info.soft_budget, team_member_budget: info.team_member_budget_table?.max_budget, team_member_budget_duration: info.team_member_budget_table?.budget_duration, + team_member_budget_limits: info.team_member_budget_limits ?? [], team_member_tpm_limit: info.team_member_budget_table?.tpm_limit, team_member_rpm_limit: info.team_member_budget_table?.rpm_limit, }} diff --git a/ui/litellm-dashboard/src/components/team/memberBudgetWindowSchema.ts b/ui/litellm-dashboard/src/components/team/memberBudgetWindowSchema.ts index 79b2927c974..4a9732ecdbc 100644 --- a/ui/litellm-dashboard/src/components/team/memberBudgetWindowSchema.ts +++ b/ui/litellm-dashboard/src/components/team/memberBudgetWindowSchema.ts @@ -24,7 +24,7 @@ export const memberBudgetWindowsSchema = budgetWindowEntriesSchema.superRefine(( }); }); -export const isMemberBudgetWindows = (value: unknown): value is BudgetWindowEntry[] => +export const isMemberBudgetWindows = (value: unknown): value is MemberBudgetWindow[] => budgetWindowEntriesSchema.safeParse(value).success; export type MemberBudgetWindow = z.infer[number]; diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts index 731755959e1..434e80db12c 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import type { MemberBudgetWindow } from "./memberBudgetWindowSchema"; import { type TeamAdminMetadataContext, @@ -365,3 +366,21 @@ describe("teamAdminDeclaredFields", () => { ]); }); }); + +describe("team admin budget windows", () => { + const windows: MemberBudgetWindow[] = [{ budget_duration: "24h", max_budget: 20 }]; + it("should send configured windows only when the proxy permits editing them", () => { + const values = { team_member_budget_limits: windows }; + const initial = { team_member_budget_limits: [] }; + expect(teamAdminSettingsChanges(values, initial, new Set(["team_member_budget_limits"]))).toEqual(values); + expect(teamAdminSettingsChanges(values, initial, new Set())).toEqual({}); + }); + it("should clear configured defaults without resending unchanged windows", () => { + const initial = { team_member_budget_limits: windows }; + const permitted = new Set(["team_member_budget_limits"]); + expect(teamAdminSettingsChanges({ team_member_budget_limits: [] }, initial, permitted)).toEqual({ + team_member_budget_limits: [], + }); + expect(teamAdminSettingsChanges(initial, initial, permitted)).toEqual({}); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index bf8d56b6110..8d4d0bf13bb 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -1,4 +1,5 @@ import { z } from "zod"; +import type { MemberBudgetWindow } from "./memberBudgetWindowSchema"; import { numberOrNull } from "@/lib/forms/numberOrNull"; import { @@ -93,6 +94,7 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([ ["soft_budget", "Soft Budget (USD)"], ["team_member_budget", "Member Default Budget (USD)"], ["team_member_budget_duration", "Member Default Budget Duration"], + ["team_member_budget_limits", "Default member budget windows"], ["team_member_tpm_limit", "Member Default TPM Limit"], ["team_member_rpm_limit", "Member Default RPM Limit"], [TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION, "Raise the team's max budget"], @@ -105,11 +107,11 @@ export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_L export type TeamAdminSettingsValues = { readonly [F in TeamAdminNumericSettingsField]?: string | number | null } & { readonly [F in TeamAdminDurationSettingsField]?: string | null; -}; +} & { readonly team_member_budget_limits?: MemberBudgetWindow[] }; export type TeamAdminSettingsChanges = { readonly [F in TeamAdminNumericSettingsField]?: number | null } & { readonly [F in TeamAdminDurationSettingsField]?: string | null; -} & { readonly metadata?: Record }; +} & { readonly metadata?: Record; readonly team_member_budget_limits?: MemberBudgetWindow[] }; const durationOrNull = (value: string | number | null | undefined): string | null => typeof value === "string" && value !== "" ? value : null; @@ -132,13 +134,19 @@ export const teamAdminSettingsChanges = ( values: TeamAdminSettingsValues, initialValues: TeamAdminSettingsValues, editableFields: ReadonlySet, -): TeamAdminSettingsChanges => - Object.fromEntries( +): TeamAdminSettingsChanges => { + const changes = Object.fromEntries( TEAM_ADMIN_SETTINGS_FIELDS.flatMap((field) => { const change = editableFields.has(field) ? changedValue(field, values[field], initialValues[field]) : null; return change === null ? [] : [change]; }), ); + if (!editableFields.has("team_member_budget_limits")) return changes; + const windows = values.team_member_budget_limits; + if (windows === undefined) return changes; + if (JSON.stringify(windows) === JSON.stringify(initialValues.team_member_budget_limits ?? [])) return changes; + return { ...changes, team_member_budget_limits: windows }; +}; export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess => { const parsed = callerEditAccessSchema.safeParse(callerEditAccess); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e80df24a818..3efa1b27b6f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8357,7 +8357,7 @@ export interface paths { * Example: * ```bash * curl --location 'http://0.0.0.0:4000/key/block' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ - * "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + * "key": "" * }' * ``` * @@ -8916,7 +8916,7 @@ export interface paths { * Example: * ```bash * curl --location 'http://0.0.0.0:4000/key/unblock' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ - * "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + * "key": "" * }' * ``` * @@ -34680,6 +34680,8 @@ export interface components { team_alias?: string | null; /** Team Id */ team_id: string; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Permissions */ team_member_permissions?: string[] | null; /** Tpd Limit */ @@ -36196,6 +36198,8 @@ export interface components { LiteLLM_TeamMembership: { /** Budget Id */ budget_id?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Litellm Budget Table */ litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTableFull"] | components["schemas"]["LiteLLM_BudgetTable"] | null; /** @@ -36303,6 +36307,8 @@ export interface components { team_alias?: string | null; /** Team Id */ team_id: string; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Permissions */ team_member_permissions?: string[] | null; /** Tpd Limit */ @@ -39677,6 +39683,8 @@ export interface components { team_member_budget?: number | null; /** Team Member Budget Duration */ team_member_budget_duration?: string | null; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Key Duration */ team_member_key_duration?: string | null; /** Team Member Model Max Budget */ @@ -40967,6 +40975,8 @@ export interface components { team_member_budget?: number | null; /** Team Member Budget Duration */ team_member_budget_duration?: string | null; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Key Duration */ team_member_key_duration?: string | null; /** Team Member Model Max Budget */ @@ -46738,6 +46748,8 @@ export interface components { team_alias?: string | null; /** Team Id */ team_id: string; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Permissions */ team_member_permissions?: string[] | null; /** Tpd Limit */ @@ -46898,6 +46910,8 @@ export interface components { team_alias?: string | null; /** Team Id */ team_id: string; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Permissions */ team_member_permissions?: string[] | null; /** Tpd Limit */ @@ -46964,6 +46978,8 @@ export interface components { * @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets. */ budget_duration?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** * Max Budget In Team * @description Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits @@ -47002,6 +47018,11 @@ export interface components { /** User Id */ user_id?: string | null; }; + /** + * TeamMemberBudgetDuration + * @enum {string} + */ + TeamMemberBudgetDuration: "1h" | "24h" | "7d" | "30d"; /** * TeamMemberBudgetPatch * @description One member's per-member limits, merge-patch style: a field left out of the row is @@ -47053,6 +47074,12 @@ export interface components { /** User Id */ user_id?: string | null; }; + /** TeamMemberBudgetWindow */ + TeamMemberBudgetWindow: { + budget_duration: components["schemas"]["TeamMemberBudgetDuration"]; + /** Max Budget */ + max_budget: number; + }; /** TeamMemberDeleteRequest */ TeamMemberDeleteRequest: { /** Team Id */ @@ -47083,6 +47110,8 @@ export interface components { TeamMemberInfoResponse: { /** Budget Id */ budget_id?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Litellm Budget Table */ litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTableFull"] | components["schemas"]["LiteLLM_BudgetTable"] | null; /** Role */ @@ -47144,6 +47173,8 @@ export interface components { * @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets. */ budget_duration?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Max Budget In Team */ max_budget_in_team?: number | null; /** @@ -47188,6 +47219,8 @@ export interface components { allowed_models?: string[] | null; /** Budget Duration */ budget_duration?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Max Budget In Team */ max_budget_in_team?: number | null; /** Model Max Budget */ @@ -49284,6 +49317,8 @@ export interface components { team_member_budget?: number | null; /** Team Member Budget Duration */ team_member_budget_duration?: string | null; + /** Team Member Budget Limits */ + team_member_budget_limits?: components["schemas"]["TeamMemberBudgetWindow"][] | null; /** Team Member Key Duration */ team_member_key_duration?: string | null; /** Team Member Model Max Budget */