diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260929000000_add_team_membership_self_max_budget/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260929000000_add_team_membership_self_max_budget/migration.sql new file mode 100644 index 00000000000..8e0304074b2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260929000000_add_team_membership_self_max_budget/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN IF NOT EXISTS "self_max_budget" DOUBLE PRECISION; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f29caa9ceb7..e647db2cc99 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -788,6 +788,7 @@ model LiteLLM_TeamMembership { spend Float @default(0.0) total_spend Float @default(0.0) budget_id String? + self_max_budget Float? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) @@id([user_id, team_id]) } diff --git a/litellm/models/team_membership.py b/litellm/models/team_membership.py index 0ffe8f8dbcc..8760791cf11 100644 --- a/litellm/models/team_membership.py +++ b/litellm/models/team_membership.py @@ -15,8 +15,19 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): budget_id: str | None = None spend: float | None = 0.0 total_spend: float | None = 0.0 + self_max_budget: float | None = None litellm_budget_table: LiteLLM_BudgetTableFull | LiteLLM_BudgetTable | None = None + def self_cap_binds(self, admin_budget: float | None) -> bool: + if self.self_max_budget is None: + return False + return admin_budget is None or self.self_max_budget < admin_budget + + def capped_budget(self, admin_budget: float | None) -> float | None: + if self.self_cap_binds(admin_budget): + return self.self_max_budget + return admin_budget + 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/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3ec430332ee..d0b43a69ff8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -14,6 +14,7 @@ import math import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from datetime import datetime from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias @@ -5849,31 +5850,19 @@ async def _check_team_member_budget( ) loaded_membership = team_membership - # Per-member override wins; otherwise fall back to the team-level - # default configured via team.metadata["team_member_budget_id"]. - team_member_budget: float | None = None - member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None now: Final = get_utc_datetime() - if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) - else: - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - if isinstance(default_budget_id, str): - default_budget: Final = await get_team_member_default_budget( - budget_id=default_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - # Treat 0 on the team default as "no cap". - # Per-member rows still respect 0 as an explicit admin disable. - if ( - default_budget is not None - and default_budget.max_budget is not None - and default_budget.max_budget > 0 - ): - team_member_budget = default_budget.max_budget + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) + admin_budget: Final = await get_team_member_admin_budget( + team_object=team_object, + team_membership=loaded_membership, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + now=now, + ) + + self_binds: Final = loaded_membership is not None and loaded_membership.self_cap_binds(admin_budget) + team_member_budget: Final = ( + loaded_membership.capped_budget(admin_budget) if loaded_membership is not None else admin_budget + ) if team_member_budget is not None: team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 @@ -5890,19 +5879,34 @@ async def _check_team_member_budget( if not math.isfinite(team_member_budget): return - _team_member_max_budget_alert_check( - team_id=team_object.team_id, - team_alias=team_object.team_alias, - team_metadata=team_object.metadata, - organization_id=team_object.organization_id, - user_id=valid_token.user_id, - user_email=user_object.user_email if user_object is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) + if admin_budget is not None and math.isfinite(admin_budget): + _team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=team_object.metadata, + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=team_member_spend, + max_budget=admin_budget, + ) if team_member_spend >= team_member_budget: + if self_binds: + raise litellm.BudgetExceededError( + current_cost=team_member_spend, + max_budget=team_member_budget, + message=( + f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} " + f"Current cost: {team_member_spend}, Max budget: {team_member_budget}. " + "This is the personal spend cap you set for yourself in this team; raise or clear " + f"it with PATCH /team/{team_object.team_id}/members/me " + "or from the My User tab of the team." + ), + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{team_object.team_id}", + ) raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, @@ -5912,6 +5916,34 @@ async def _check_team_member_budget( ) +async def get_team_member_admin_budget( + team_object: LiteLLM_TeamTable, + team_membership: LiteLLM_TeamMembership | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + now: datetime, +) -> float | None: + """The admin-set per-member budget, or None when no admin cap applies.""" + member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None + if member_budget_row is not None and member_budget_row.max_budget is not None: + return member_budget_row.effective_max_budget(now=now) + + default_budget_id: Final = ( + team_object.metadata or {} # mutable-ok: read-only empty fallback + ).get("team_member_budget_id") + if isinstance(default_budget_id, str): + default_budget: Final = await get_team_member_default_budget( + budget_id=default_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if default_budget is not None and default_budget.max_budget is not None and default_budget.max_budget > 0: + return default_budget.max_budget + ( + member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 + ) + return None + + async def _check_team_member_model_access( model: str | list[str], team_object: LiteLLM_TeamTable, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6cec3e714ec..5f0ee65d8c5 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -104,6 +104,7 @@ from litellm.proxy.auth.auth_checks import ( delete_cache_team_object, get_jwt_key_mapping_cache_keys_for_tokens, get_org_object, + get_team_member_admin_budget, get_team_membership, get_team_object, get_user_object, @@ -208,13 +209,16 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamListItem, TeamListResponse, TeamMemberAddResult, + TeamMemberEffectiveBudgetSource, TeamMemberInfoResponse, + TeamMemberSelfBudgetUpdateRequest, TeamMetadataSchemaResponse, TeamUserSpendResponse, TeamUserSpendRow, UpdateTeamMemberPermissionsRequest, ) from litellm.types.utils import BudgetConfig +from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import Prisma @@ -5073,17 +5077,7 @@ async def team_member_me( user_api_key_cache=user_api_key_cache, ) - caller_user_email: Final = user_api_key_dict.user_email - member_role: str | None = None - for m in team_table.members_with_roles: - # Match by user_id when present, else fall back to email — members - # added by email may have user_id=None on the stored entry. - if (m.user_id is not None and m.user_id == caller_user_id) or ( - m.user_email is not None and caller_user_email is not None and m.user_email == caller_user_email - ): - member_role = m.role - break - + member_role: Final = _caller_team_member_role(team_table=team_table, user_api_key_dict=user_api_key_dict) if member_role is None: # Caller is not a member of this team. Even proxy admins get 404 here — # they can use /team/info to view all members; "me" only resolves for @@ -5108,6 +5102,71 @@ async def team_member_me( ) user_email: Final = getattr(user_row, "user_email", None) if user_row is not None else None + return await _build_team_member_info_response( + team_id=team_id, + team_table=team_table, + member_role=member_role, + caller_user_id=caller_user_id, + user_email=user_email, + membership=membership, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + +def _caller_team_member_role(team_table: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth) -> str | None: + caller_user_id: Final = user_api_key_dict.user_id + caller_user_email: Final = user_api_key_dict.user_email + for m in team_table.members_with_roles: + if (m.user_id is not None and m.user_id == caller_user_id) or ( + m.user_email is not None and caller_user_email is not None and m.user_email == caller_user_email + ): + return m.role + return None + + +def _member_effective_budget_source( + membership: LiteLLM_TeamMembership | None, + admin_budget: float | None, + team_table: LiteLLM_TeamTable, +) -> TeamMemberEffectiveBudgetSource: + if membership is not None and membership.self_cap_binds(admin_budget): + return "self" + member_budget_row: Final = membership.litellm_budget_table if membership is not None else None + if membership is not None and member_budget_row is not None and member_budget_row.max_budget is not None: + default_budget_id: Final = ( + team_table.metadata or {} # mutable-ok: read-only empty fallback + ).get("team_member_budget_id") + return _member_budget_source( + membership.budget_id, + default_budget_id if isinstance(default_budget_id, str) else None, + ) + return "team_default" if admin_budget is not None else "none" + + +async def _build_team_member_info_response( + team_id: str, + team_table: LiteLLM_TeamTable, + member_role: str, + caller_user_id: str, + user_email: str | None, + membership: LiteLLM_TeamMembership | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> TeamMemberInfoResponse: + admin_budget: Final = await get_team_member_admin_budget( + team_object=team_table, + team_membership=membership, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + now=get_utc_datetime(), + ) + effective_budget: Final = membership.capped_budget(admin_budget) if membership is not None else admin_budget + budget_source: Final = ( + _member_effective_budget_source(membership=membership, admin_budget=admin_budget, team_table=team_table) + if effective_budget is not None + else "none" + ) if membership is None: # Member is in members_with_roles but has no membership row yet # (no per-member budget/limits configured). Return defaults. @@ -5120,9 +5179,11 @@ async def team_member_me( spend=0.0, total_spend=0.0, budget_id=None, + self_max_budget=None, litellm_budget_table=None, + effective_budget=effective_budget, + budget_source=budget_source, ) - return TeamMemberInfoResponse( user_id=caller_user_id, team_id=team_id, @@ -5132,7 +5193,206 @@ async def team_member_me( spend=membership.spend, total_spend=membership.total_spend, budget_id=membership.budget_id, + self_max_budget=membership.self_max_budget, litellm_budget_table=membership.litellm_budget_table, + effective_budget=effective_budget, + budget_source=budget_source, + ) + + +@router.patch( + "/team/{team_id}/members/me", + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence + dependencies=(Depends(user_api_key_auth),), + response_model=TeamMemberInfoResponse, +) +@management_endpoint_wrapper +async def team_member_me_update_self_budget( + http_request: Request, + team_id: str, + data: TeamMemberSelfBudgetUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> TeamMemberInfoResponse: + """ + Set or clear the caller's own personal spend cap within the team. + + The self cap can only lower the member's effective budget below their + team allocation, never raise it: values above the current admin-set + allocation are rejected, and the cap binds only while it stays below + that allocation. Send `self_max_budget: null` to clear it. + + Returns 404 if the caller is not a member of the team. + + ``` + curl --location --request PATCH 'http://localhost:4000/team/your_team_id/members/me' \ + --header 'Authorization: Bearer your_api_key_here' \ + --header 'Content-Type: application/json' \ + --data '{"self_max_budget": 80.0}' + ``` + """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + is_audit_logging_enabled, + ) + from litellm.proxy.proxy_server import ( + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape + "error": "Database not connected. Connect a database to your proxy - " + "https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + }, + ) + + caller_user_id: Final = user_api_key_dict.user_id + if caller_user_id is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape + "error": "API key has no associated user_id; cannot resolve 'me' for team membership." + }, + ) + + team_table: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + member_role: Final = _caller_team_member_role(team_table=team_table, user_api_key_dict=user_api_key_dict) + if member_role is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape + "error": f"User user_id={caller_user_id} is not a member of team_id={team_id}." + }, + ) + + membership_row: Final = await _team_membership_db(prisma_client).find_unique( + where={ # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": { # mutable-ok: Prisma query filters are dict-shaped + "user_id": caller_user_id, + "team_id": team_id, + } + }, + include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument + ) + membership: Final = ( + LiteLLM_TeamMembership.model_validate(membership_row.model_dump()) if membership_row is not None else None + ) + + admin_budget: Final = await get_team_member_admin_budget( + team_object=team_table, + team_membership=membership, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + now=get_utc_datetime(), + ) + + new_self_max_budget: Final = data.self_max_budget + if new_self_max_budget is not None and admin_budget is not None and new_self_max_budget > admin_budget: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape + "error": ( + f"self_max_budget {new_self_max_budget} exceeds your current team allocation of {admin_budget}." + ) + }, + ) + + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": { # mutable-ok: Prisma query filters are dict-shaped + "user_id": caller_user_id, + "team_id": team_id, + } + } + if membership_row is not None or new_self_max_budget is not None: + await _team_membership_db(prisma_client).upsert( + where=membership_where, + data={ # mutable-ok: prisma client requires a plain dict data= argument + "create": { # mutable-ok: same prisma data= argument + "user_id": caller_user_id, + "team_id": team_id, + "self_max_budget": new_self_max_budget, + }, + "update": {"self_max_budget": new_self_max_budget}, # mutable-ok: same prisma data= argument + }, + ) + await invalidate_team_member_spend_state( + user_id=caller_user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + ) + + if is_audit_logging_enabled(): + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=team_id, + action="updated", + updated_values=json.dumps( + { # mutable-ok: the audit-log JSON column rejects a top-level array, so this value must be an object + "user_id": caller_user_id, + "self_max_budget": new_self_max_budget, + } + ), + before_value=json.dumps( + { # mutable-ok: the audit-log JSON column rejects a top-level array, so this value must be an object + "user_id": caller_user_id, + "self_max_budget": membership.self_max_budget if membership is not None else None, + } + ), + ) + ) + ) + + updated_membership: Final = ( + membership.model_copy( + update={ # mutable-ok: pydantic update payload + "self_max_budget": new_self_max_budget, + } + ) + if membership is not None + else ( + LiteLLM_TeamMembership(user_id=caller_user_id, team_id=team_id, self_max_budget=new_self_max_budget) + if new_self_max_budget is not None + else None + ) + ) + + user_row: Final = await get_user_object( + user_id=caller_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + user_email: Final = user_row.user_email if user_row is not None else None + + return await _build_team_member_info_response( + team_id=team_id, + team_table=team_table, + member_role=member_role, + caller_user_id=caller_user_id, + user_email=user_email, + membership=updated_membership, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f29caa9ceb7..e647db2cc99 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -788,6 +788,7 @@ model LiteLLM_TeamMembership { spend Float @default(0.0) total_spend Float @default(0.0) budget_id String? + self_max_budget Float? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) @@id([user_id, team_id]) } diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index c094e91c6c0..87cfe0c7586 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -753,9 +753,9 @@ async def _get_team_member_budget_counter( member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None now: Final = datetime.now(timezone.utc) - team_member_budget: float | None = None + admin_budget: float | None = None if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) + admin_budget = member_budget_row.effective_max_budget(now=now) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): @@ -764,10 +764,14 @@ async def _get_team_member_budget_counter( ) default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) if default_cap is not None and default_cap > 0: - team_member_budget = default_cap + ( + admin_budget = default_cap + ( member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 ) + team_member_budget: Final = ( + team_membership.capped_budget(admin_budget) if team_membership is not None else admin_budget + ) + if team_member_budget is None or team_member_budget <= 0: return None diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 4524c47ec38..d72d101221d 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,4 +1,4 @@ -from typing import Any, Final, Literal +from typing import Annotated, Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -207,12 +207,27 @@ class BulkTeamMemberBudgetUpdateResponse(ResourceResponse[tuple[TeamMemberBudget """`{data: [...]}` with one `TeamMemberBudgetUpdateResult` per requested member, in request order.""" +TeamMemberEffectiveBudgetSource: TypeAlias = Literal["team_default", "custom", "self", "none"] + + class TeamMemberInfoResponse(LiteLLM_TeamMembership): """Response for GET /team/{team_id}/members/me — caller's own membership row.""" role: str | None = None user_email: str | None = None team_alias: str | None = None + effective_budget: float | None = None + budget_source: TeamMemberEffectiveBudgetSource = "none" + + +class TeamMemberSelfBudgetUpdateRequest(BaseModel): + """Body of `PATCH /team/{team_id}/members/me`. `self_max_budget` is + required: a non-negative number sets the caller's personal cap, null + clears it.""" + + model_config = ConfigDict(extra="forbid") + + self_max_budget: Annotated[float, Field(ge=0, allow_inf_nan=False, strict=True)] | None class TeamMetadataFieldSchema(BaseModel): diff --git a/schema.prisma b/schema.prisma index f29caa9ceb7..e647db2cc99 100644 --- a/schema.prisma +++ b/schema.prisma @@ -788,6 +788,7 @@ model LiteLLM_TeamMembership { spend Float @default(0.0) total_spend Float @default(0.0) budget_id String? + self_max_budget Float? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) @@id([user_id, team_id]) } diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 99ed82dcad8..cf946e1a18b 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -47,6 +47,7 @@ GET /search_tools/ui/available_providers GET /team/available GET /team/metadata_schema GET /team/{team_id}/members/me +PATCH /team/{team_id}/members/me GET /user/available_users # Imperative one-shot operations: bulk edits, rotation, health probes, test hooks, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 353249dddf0..4898d481021 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7927,13 +7927,12 @@ async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_ async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request + from litellm.proxy._types import LiteLLM_TeamMembership from litellm.proxy.auth.auth_checks import common_checks team = LiteLLM_TeamTable(team_id="t-once") token = UserAPIKeyAuth(token="k-once", user_id="u-once", team_id="t-once", models=["gpt-4o-mini"]) - membership = MagicMock() - membership.litellm_budget_table = None - membership.spend = 0.0 + membership = LiteLLM_TeamMembership(user_id="u-once", team_id="t-once", spend=0.0) with ( patch( # test-quality-ok: common_checks imports prisma_client from proxy_server @@ -10276,3 +10275,241 @@ async def test_authoritative_group_grants_propagate_policy_outages( await _get_agent_ids_from_access_groups(["group"], check_db_only=True) else: assert await _get_agent_ids_from_access_groups(["group"]) == [] + + +def _team_member_check_setup( + membership, + spend: float, + team_metadata: dict[str, str] | None = None, +): + """Return the shared fixtures for a _check_team_member_budget call.""" + from litellm.caching.dual_cache import DualCache + + team_object = LiteLLM_TeamTable(team_id=membership.team_id, metadata=team_metadata or {}) + user_object = LiteLLM_UserTable(user_id=membership.user_id) + valid_token = UserAPIKeyAuth(token="test-token", user_id=membership.user_id, team_id=membership.team_id) + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + if counter_key == f"spend:team_member:{membership.user_id}:{membership.team_id}": + return spend + return fallback_spend + + ctx = ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=membership, + ), + ) + return team_object, user_object, valid_token, ctx, DualCache() + + +@pytest.mark.asyncio +async def test_team_member_budget_self_cap_binds_below_admin_budget(): + """A self cap below the member-row admin budget enforces the self cap, + and the raised error points the member at the self-cap route.""" + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + binding_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + self_max_budget=80.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + team_object, user_object, valid_token, ctx, cache = _team_member_check_setup(binding_membership, 85.0) + with ctx[0], ctx[1]: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + assert exc_info.value.max_budget == 80.0 + assert "personal spend cap" in str(exc_info.value.message) + + loose_membership = binding_membership.model_copy(update={"self_max_budget": 90.0}) + team_object, user_object, valid_token, ctx, cache = _team_member_check_setup(loose_membership, 85.0) + with ctx[0], ctx[1]: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + + tied_membership = binding_membership.model_copy(update={"self_max_budget": 100.0}) + team_object, user_object, valid_token, ctx, cache = _team_member_check_setup(tied_membership, 100.0) + with ctx[0], ctx[1]: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + assert "personal spend cap" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_team_member_budget_self_cap_against_team_default_with_temp_increase(): + """The self cap is measured against the live admin allocation: the team + default plus an active temp increase. Above it the admin side binds and + no self-cap error fires; below it the self cap does.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + cache = DualCache() + await cache.async_set_cache( + key="team_member_default_budget:default-budget-1", + value=LiteLLM_BudgetTable(budget_id="default-budget-1", max_budget=100.0), + ) + budget_row = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=50.0, + temp_budget_expiry=datetime.now(timezone.utc) + timedelta(hours=1), + ) + + async def run(self_cap, spend): + membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + self_max_budget=self_cap, + budget_id="budget-1", + litellm_budget_table=budget_row, + ) + team_object = LiteLLM_TeamTable(team_id="test-team", metadata={"team_member_budget_id": "default-budget-1"}) + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + if counter_key == "spend:team_member:test-user:test-team": + return spend + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=membership, + ), + ): + await _check_team_member_budget( + team_object=team_object, + user_object=LiteLLM_UserTable(user_id="test-user"), + valid_token=UserAPIKeyAuth(token="test-token", user_id="test-user", team_id="test-team"), + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await run(self_cap=120.0, spend=130.0) + assert exc_info.value.max_budget == 120.0 + assert "personal spend cap" in str(exc_info.value.message) + + await run(self_cap=None, spend=130.0) + + await run(self_cap=200.0, spend=130.0) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await run(self_cap=200.0, spend=160.0) + assert exc_info.value.max_budget == 150.0 + assert "personal spend cap" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_team_member_budget_self_cap_enforced_without_any_admin_budget(): + """A self cap alone, with no member-row budget and no team default, is + still enforced.""" + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + self_max_budget=50.0, + budget_id=None, + litellm_budget_table=None, + ) + team_object, user_object, valid_token, ctx, cache = _team_member_check_setup(membership, 60.0) + with ctx[0], ctx[1]: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + assert exc_info.value.max_budget == 50.0 + assert "personal spend cap" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_team_member_budget_alert_uses_admin_budget_never_self_cap(): + """The team-admin alert must be computed against the admin allocation: + a binding self cap below it must not fire it, and a self cap above the + admin threshold must not suppress it.""" + from litellm.proxy._types import LiteLLM_TeamMembership + + captured: list[tuple[str, CallInfo]] = [] + + class RecordingProxyLogging: + async def budget_alerts(self, type, user_info): + captured.append((type, user_info)) + + metadata = {"team_member_max_budget_alert_emails": {"80": ["admin@co.com"]}} + + binding_membership = LiteLLM_TeamMembership( + user_id="user-1", + team_id="team-1", + spend=0.0, + self_max_budget=50.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + team_object, user_object, valid_token, ctx, cache = _team_member_check_setup( + binding_membership, 60.0, team_metadata=metadata + ) + with ctx[0], ctx[1]: + with pytest.raises(litellm.BudgetExceededError): + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=RecordingProxyLogging(), + ) + await asyncio.sleep(0) + assert captured == [], captured + + loose_membership = binding_membership.model_copy(update={"self_max_budget": 95.0}) + team_object, user_object, valid_token, ctx, cache = _team_member_check_setup( + loose_membership, 85.0, team_metadata=metadata + ) + with ctx[0], ctx[1]: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=RecordingProxyLogging(), + ) + await asyncio.sleep(0) + assert len(captured) == 1, captured + assert captured[0][0] == "max_budget_alert" + assert captured[0][1].max_budget == 100.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a53894fcd1b..4cf0ad32c7c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13768,11 +13768,11 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp _wire_member_delete_tx(mock_prisma_client) with ( - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.team_info", AsyncMock(side_effect=[_team_info_as_read_from_db("user"), _team_info_as_read_from_db("admin")]), ), - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", AsyncMock(), ), @@ -13818,7 +13818,7 @@ def _roster_writer(team_row: LiteLLM_TeamTable): def _member_update_patches(team_snapshot: LiteLLM_TeamTable): return ( - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.team_info", AsyncMock( return_value={ @@ -13829,7 +13829,7 @@ def _member_update_patches(team_snapshot: LiteLLM_TeamTable): } ), ), - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", AsyncMock(), ), @@ -15195,7 +15195,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), ): @@ -15228,7 +15228,7 @@ async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), ): @@ -15251,7 +15251,7 @@ async def test_reset_team_member_spend_fn_team_not_found(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})), ): @@ -15276,7 +15276,7 @@ async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), ): @@ -15304,7 +15304,7 @@ async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monk monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1") - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock( return_value=LiteLLM_TeamTable( @@ -15337,7 +15337,7 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), ): @@ -15381,7 +15381,7 @@ async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")), ): @@ -15422,7 +15422,7 @@ async def test_reset_team_member_budget_fn_detaches_member_when_team_has_no_usab monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=team_obj), ): @@ -15448,7 +15448,7 @@ async def test_reset_team_member_budget_fn_membership_not_found(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")), ): @@ -15468,7 +15468,7 @@ async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + with patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), ): @@ -15499,7 +15499,7 @@ async def _team_info_budget_sources( mock_prisma.get_data = AsyncMock(return_value=[]) with ( - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.proxy_server.prisma_client", mock_prisma ), patch.object( # test-quality-ok: membership lookup is a module-level DB query with no injection point @@ -15605,11 +15605,11 @@ async def test_team_member_update_invalidates_team_member_spend_state_when_budge mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) with ( - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.team_info", AsyncMock(return_value=team_info_response), ), - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", AsyncMock(), ), @@ -15658,11 +15658,11 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent( mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) with ( - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints.team_info", AsyncMock(return_value=team_info_response), ), - patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + patch( # test-quality-ok: no live DB in this unit test "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", AsyncMock(), ), @@ -16889,3 +16889,605 @@ def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hi assert response.status_code == 503, response.text assert response.json() == _DB_OUTAGE_503_BODY + + +def _member_self_cap_auth(user_id: str = "member-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-member", user_id=user_id) + + +def _team_with_member(team_id: str, user_id: str, role: str = "user") -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, members_with_roles=[Member(user_id=user_id, role=role)]) + + +def _wire_self_cap_route(monkeypatch, team_obj, membership_row): + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.upsert = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_team_member_me_patch_sets_own_self_cap(monkeypatch): + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + audit_logger = _wire_audit_log_callback(monkeypatch) + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + spend=10.0, + self_max_budget=60.0, + budget_id="b1", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=100.0), + ) + mock_prisma_client = _wire_self_cap_route(monkeypatch, _team_with_member("team-1", "member-1"), membership_row) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_member("team-1", "member-1")), + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: the cache fan-out needs a redis bus this test does not have + "litellm.proxy.management_endpoints.team_endpoints.invalidate_team_member_spend_state", + new_callable=AsyncMock, + ) as mock_invalidate, + ): + response = await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=80.0), + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + await _settle_audit_log_tasks() + + mock_prisma_client.db.litellm_teammembership.upsert.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={ + "create": {"user_id": "member-1", "team_id": "team-1", "self_max_budget": 80.0}, + "update": {"self_max_budget": 80.0}, + }, + ) + mock_invalidate.assert_awaited_once() + assert mock_invalidate.call_args.kwargs["user_id"] == "member-1" + assert mock_invalidate.call_args.kwargs["team_id"] == "team-1" + assert response.self_max_budget == 80.0 + assert response.effective_budget == 80.0 + assert response.budget_source == "self" + + updated_events = _team_roster_events(audit_logger, "updated") + assert len(updated_events) == 1 + assert updated_events[0]["object_id"] == "team-1" + assert json.loads(updated_events[0]["updated_values"]) == {"user_id": "member-1", "self_max_budget": 80.0} + assert json.loads(updated_events[0]["before_value"]) == {"user_id": "member-1", "self_max_budget": 60.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("self_cap, expect_ok", [(120.0, False), (100.0, True)]) +async def test_team_member_me_patch_caps_at_admin_allocation(monkeypatch, self_cap, expect_ok): + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=100.0), + ) + mock_prisma_client = _wire_self_cap_route(monkeypatch, _team_with_member("team-1", "member-1"), membership_row) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_member("team-1", "member-1")), + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: the cache fan-out needs a redis bus this test does not have + "litellm.proxy.management_endpoints.team_endpoints.invalidate_team_member_spend_state", + new_callable=AsyncMock, + ), + ): + if not expect_ok: + with pytest.raises(HTTPException) as exc: + await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=self_cap), + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + assert exc.value.status_code == 400 + assert "exceeds your current team allocation" in str(exc.value.detail) + mock_prisma_client.db.litellm_teammembership.upsert.assert_not_awaited() + return + response = await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=self_cap), + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + assert response.self_max_budget == self_cap + + +@pytest.mark.asyncio +async def test_team_member_me_patch_honours_team_default_plus_temp_increase(monkeypatch): + """The allocation ceiling includes a live team default and an active temp + increase, so a self cap between the default and the inflated total is accepted.""" + from datetime import timedelta + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="team_member_default_budget:default-b1", + value=LiteLLM_BudgetTable(budget_id="default-b1", max_budget=100.0), + ) + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="b1", + litellm_budget_table=LiteLLM_BudgetTable( + budget_id="b1", + max_budget=None, + temp_budget_increase=50.0, + temp_budget_expiry=datetime.now(timezone.utc) + timedelta(hours=1), + ), + ) + team_obj = LiteLLM_TeamTable( + team_id="team-1", + metadata={"team_member_budget_id": "default-b1"}, + members_with_roles=[Member(user_id="member-1", role="user")], + ) + _wire_self_cap_route(monkeypatch, team_obj, membership_row) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=team_obj), + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: the cache fan-out needs a redis bus this test does not have + "litellm.proxy.management_endpoints.team_endpoints.invalidate_team_member_spend_state", + new_callable=AsyncMock, + ), + ): + response = await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=140.0), + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + assert response.effective_budget == 140.0 + assert response.budget_source == "self" + + +@pytest.mark.asyncio +async def test_team_member_me_patch_null_clears_self_cap(monkeypatch): + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + self_max_budget=80.0, + budget_id="b1", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=100.0), + ) + mock_prisma_client = _wire_self_cap_route(monkeypatch, _team_with_member("team-1", "member-1"), membership_row) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_member("team-1", "member-1")), + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: the cache fan-out needs a redis bus this test does not have + "litellm.proxy.management_endpoints.team_endpoints.invalidate_team_member_spend_state", + new_callable=AsyncMock, + ), + ): + response = await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=None), + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + + update_data = mock_prisma_client.db.litellm_teammembership.upsert.call_args.kwargs["data"]["update"] + assert update_data == {"self_max_budget": None} + assert response.self_max_budget is None + assert response.effective_budget == 100.0 + assert response.budget_source == "custom" + + +@pytest.mark.asyncio +async def test_team_member_me_patch_null_with_no_membership_row_writes_nothing(monkeypatch): + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + audit_logger = _wire_audit_log_callback(monkeypatch) + mock_prisma_client = _wire_self_cap_route(monkeypatch, _team_with_member("team-1", "member-1"), None) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_member("team-1", "member-1")), + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: the cache fan-out needs a redis bus this test does not have + "litellm.proxy.management_endpoints.team_endpoints.invalidate_team_member_spend_state", + new_callable=AsyncMock, + ) as mock_invalidate, + ): + response = await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=None), + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + + await _settle_audit_log_tasks() + + mock_prisma_client.db.litellm_teammembership.upsert.assert_not_awaited() + mock_invalidate.assert_not_awaited() + assert response.self_max_budget is None + assert response.budget_source == "none" + assert _team_roster_events(audit_logger, "updated") == [] + + +@pytest.mark.asyncio +async def test_team_member_me_patch_non_member_gets_404(monkeypatch): + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + mock_prisma_client = _wire_self_cap_route( + monkeypatch, LiteLLM_TeamTable(team_id="team-1", members_with_roles=[]), None + ) + + with patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), + ): + with pytest.raises(HTTPException) as exc: + await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=50.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + mock_prisma_client.db.litellm_teammembership.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_member_me_patch_key_without_user_id_gets_400(monkeypatch): + from litellm.proxy.management_endpoints.team_endpoints import team_member_me_update_self_budget + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + mock_prisma_client = _wire_self_cap_route(monkeypatch, _team_with_member("team-1", "member-1"), None) + + with pytest.raises(HTTPException) as exc: + await team_member_me_update_self_budget( + http_request=MagicMock(), + team_id="team-1", + data=TeamMemberSelfBudgetUpdateRequest(self_max_budget=50.0), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, api_key="sk-team"), + ) + assert exc.value.status_code == 400 + mock_prisma_client.db.litellm_teammembership.upsert.assert_not_awaited() + + +@pytest.mark.parametrize( + "body", + [ + {"self_max_budget": 50, "user_id": "other"}, + {"self_max_budget": True}, + {"self_max_budget": "50"}, + {"self_max_budget": -1}, + {"self_max_budget": float("inf")}, + {}, + ], + ids=["extra_key", "bool", "string", "negative", "inf", "missing_field"], +) +def test_team_member_self_budget_update_request_rejects_invalid_bodies(body): + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + with pytest.raises(ValidationError): + TeamMemberSelfBudgetUpdateRequest.model_validate(body) + + +@pytest.mark.parametrize("value", [80, 80.5, 0]) +def test_team_member_self_budget_update_request_accepts_numbers(value): + from litellm.types.proxy.management_endpoints.team_endpoints import TeamMemberSelfBudgetUpdateRequest + + request: Final = TeamMemberSelfBudgetUpdateRequest.model_validate({"self_max_budget": value}) + assert request.self_max_budget == float(value) + + +def test_team_member_me_patch_route_rejects_extra_body_key(): + """Through the real route: a body key that is not self_max_budget is a 422.""" + app.dependency_overrides[user_api_key_auth] = lambda: _member_self_cap_auth("member-1") + try: + response = client.patch("/team/team-1/members/me", json={"self_max_budget": 50, "user_id": "other"}) + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + assert response.status_code == 422, response.text + assert "user_id" in response.text + + +@pytest.mark.asyncio +async def test_team_member_me_get_reports_self_cap_and_source(monkeypatch): + """GET returns self_max_budget plus the effective budget and its source.""" + from litellm.proxy.management_endpoints.team_endpoints import team_member_me + + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + self_max_budget=80.0, + budget_id="b1", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=100.0), + ) + _wire_self_cap_route(monkeypatch, _team_with_member("team-1", "member-1"), membership_row) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_member("team-1", "member-1")), + ), + patch( # test-quality-ok: the membership fetch needs a live DB/cache layer this test does not have + "litellm.proxy.management_endpoints.team_endpoints.get_team_membership", + new_callable=AsyncMock, + return_value=membership_row, + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + response = await team_member_me( + http_request=MagicMock(), + team_id="team-1", + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + + assert response.self_max_budget == 80.0 + assert response.effective_budget == 80.0 + assert response.budget_source == "self" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "membership_row, expected_source, expected_budget", + [ + ( + LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="b1", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=100.0), + ), + "custom", + 100.0, + ), + (LiteLLM_TeamMembership(user_id="member-1", team_id="team-1"), "team_default", 60.0), + (None, "none", None), + ], + ids=["custom_admin_budget", "team_default", "no_membership_no_default"], +) +async def test_team_member_me_get_reports_budget_source( + monkeypatch, membership_row, expected_source, expected_budget +): + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_member_me + + team_obj = ( + LiteLLM_TeamTable( + team_id="team-1", + metadata={"team_member_budget_id": "default-b1"}, + members_with_roles=[Member(user_id="member-1", role="user")], + ) + if expected_source == "team_default" + else _team_with_member("team-1", "member-1") + ) + + cache = UserApiKeyCache() + if expected_source == "team_default": + await cache.async_set_cache( + key="team_member_default_budget:default-b1", + value=LiteLLM_BudgetTable(budget_id="default-b1", max_budget=60.0), + ) + + _wire_self_cap_route(monkeypatch, team_obj, membership_row) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=team_obj), + ), + patch( # test-quality-ok: the membership fetch needs a live DB/cache layer this test does not have + "litellm.proxy.management_endpoints.team_endpoints.get_team_membership", + new_callable=AsyncMock, + return_value=membership_row, + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + response = await team_member_me( + http_request=MagicMock(), + team_id="team-1", + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + + assert response.budget_source == expected_source + assert response.effective_budget == expected_budget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "self_max_budget, expected_source, expected_budget", + [(None, "team_default", 100.0), (40.0, "self", 40.0)], + ids=["member_row_without_max_uses_team_default", "self_cap_overrides_team_default"], +) +async def test_team_member_me_get_labels_team_default_when_member_row_has_no_max( + monkeypatch, self_max_budget, expected_source, expected_budget +): + """A member whose own budget row lacks max_budget falls back to the team + default, so the source label must say team_default, not custom.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_member_me + + team_obj = LiteLLM_TeamTable( + team_id="team-1", + metadata={"team_member_budget_id": "team-default-budget"}, + members_with_roles=[Member(user_id="member-1", role="user")], + ) + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="member-own-budget", + self_max_budget=self_max_budget, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="member-own-budget", max_budget=None), + ) + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="team_member_default_budget:team-default-budget", + value=LiteLLM_BudgetTable(budget_id="team-default-budget", max_budget=100.0), + ) + _wire_self_cap_route(monkeypatch, team_obj, membership_row) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=team_obj), + ), + patch( # test-quality-ok: the membership fetch needs a live DB/cache layer this test does not have + "litellm.proxy.management_endpoints.team_endpoints.get_team_membership", + new_callable=AsyncMock, + return_value=membership_row, + ), + patch( # test-quality-ok: same no-live-DB convention for the user email lookup + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + response = await team_member_me( + http_request=MagicMock(), + team_id="team-1", + user_api_key_dict=_member_self_cap_auth("member-1"), + ) + + assert response.effective_budget == expected_budget + assert response.budget_source == expected_source + + +@pytest.mark.asyncio +async def test_team_member_update_never_writes_self_max_budget(monkeypatch): + """/team/member_update must not touch a member's self cap: no membership + write carries self_max_budget, so a previously set cap survives.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team_row = LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[Member(user_id="admin-1", role="admin"), Member(user_id="member-1", role="user")], + ) + membership_with_cap = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", self_max_budget=42.0) + team_info_payload = { + "team_info": TeamInfoResponseObjectTeamTable( + team_id="team-1", + metadata={}, + members_with_roles=( + TeamInfoMember(user_id="admin-1", role="admin"), + TeamInfoMember(user_id="member-1", role="user"), + ), + ), + "team_memberships": [membership_with_cap], + } + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock() + mock_prisma_client.db.litellm_teammembership.upsert = AsyncMock() + mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock() + _wire_member_delete_tx(mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + membership_writes: list[dict] = [] + + def _record_write(_func): + async def _wrapped(*args, **kwargs): + membership_writes.append(dict(kwargs)) + return membership_with_cap + + return _wrapped + + with ( + patch( # test-quality-ok: no live DB in this unit test + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_payload), + ), + patch( # test-quality-ok: budget upsert asserted via membership table calls + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + new_callable=AsyncMock, + ) as mock_upsert_membership, + patch( # test-quality-ok: the cache fan-out needs a redis bus this test does not have + "litellm.proxy.management_endpoints.team_endpoints.invalidate_team_member_spend_state", + new_callable=AsyncMock, + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=90.0), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert mock_upsert_membership.await_count == 1 + budget_patch = mock_upsert_membership.call_args.kwargs["budget_patch"] + assert "self_max_budget" not in budget_patch + for table_call in ( + *mock_prisma_client.db.litellm_teammembership.update.call_args_list, + *mock_prisma_client.db.litellm_teammembership.upsert.call_args_list, + *mock_prisma_client.db.litellm_teammembership.update_many.call_args_list, + ): + payload = table_call.kwargs.get("data", {}) + assert "self_max_budget" not in str(payload) diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 9df8e6f4d67..8eba467f4f7 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -342,3 +342,87 @@ async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_call assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_team_member_reservation_counter_self_cap_binds_below_admin_budget() -> None: + user_id: Final = "member-self" + team_id: Final = "team-self" + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=0.5, + self_max_budget=80.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ), + ) + + counter: Final = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id), + team_object=LiteLLM_TeamTable(team_id=team_id), + user_object=LiteLLM_UserTable(user_id=user_id), + user_api_key_cache=cache, + ) + + assert counter is not None + assert counter.max_budget == 80.0 + assert counter.fallback_spend == 0.5 + + +@pytest.mark.asyncio +async def test_team_member_reservation_counter_admin_budget_binds_over_self_cap() -> None: + user_id: Final = "member-self2" + team_id: Final = "team-self2" + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=0.5, + self_max_budget=200.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ), + ) + + counter: Final = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id), + team_object=LiteLLM_TeamTable(team_id=team_id), + user_object=LiteLLM_UserTable(user_id=user_id), + user_api_key_cache=cache, + ) + + assert counter is not None + assert counter.max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_team_member_reservation_counter_self_cap_alone_creates_counter() -> None: + user_id: Final = "member-self3" + team_id: Final = "team-self3" + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=0.5, + self_max_budget=42.0, + budget_id=None, + litellm_budget_table=None, + ), + ) + + counter: Final = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id), + team_object=LiteLLM_TeamTable(team_id=team_id), + user_object=LiteLLM_UserTable(user_id=user_id), + user_api_key_cache=cache, + ) + + assert counter is not None + assert counter.max_budget == 42.0 + assert counter.fallback_spend == 0.5 diff --git a/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx b/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx index d9ba3d82230..98abd983b5a 100644 --- a/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx @@ -1,10 +1,12 @@ +import { fireEvent } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen } from "../../../tests/test-utils"; import MyUserTab from "./MyUserTab"; -import { useMyTeamMember } from "./useMyTeamMember"; +import { useMyTeamMember, useUpdateMySelfBudget } from "./useMyTeamMember"; vi.mock("./useMyTeamMember", () => ({ useMyTeamMember: vi.fn(), + useUpdateMySelfBudget: vi.fn(), })); describe("MyUserTab", () => { @@ -46,4 +48,92 @@ describe("MyUserTab", () => { expect(screen.getByText("model-one")).toBeInTheDocument(); expect(screen.getByText("TPM: 1,000")).toBeInTheDocument(); }); + + const memberInfo = (overrides: Record = {}) => ({ + user_id: "user-1", + team_id: "team-1", + role: "user", + spend: 20, + total_spend: 30, + litellm_budget_table: { max_budget: 100 }, + ...overrides, + }); + + const mockMutation = () => { + const mutate = vi.fn(); + vi.mocked(useUpdateMySelfBudget).mockReturnValue({ + mutate, + isPending: false, + isError: false, + reset: vi.fn(), + } as unknown as ReturnType); + return mutate; + }; + + it("shows the effective budget and its source badge from the API", () => { + vi.mocked(useMyTeamMember).mockReturnValue({ + data: memberInfo({ effective_budget: 80, self_max_budget: 80, budget_source: "self" }), + isLoading: false, + error: null, + } as ReturnType); + mockMutation(); + + renderWithProviders(); + + expect(screen.getByText("of $80.0000")).toBeInTheDocument(); + expect(screen.getByTestId("budget-source-badge")).toHaveTextContent("Set by you"); + expect(screen.getByTestId("my-limit-value")).toHaveTextContent("$80.0000"); + }); + + it("saves a typed limit through the PATCH mutation", () => { + vi.mocked(useMyTeamMember).mockReturnValue({ + data: memberInfo({ effective_budget: 100, budget_source: "custom", self_max_budget: null }), + isLoading: false, + error: null, + } as ReturnType); + const mutate = mockMutation(); + + renderWithProviders(); + expect(screen.getByTestId("my-limit-value")).toHaveTextContent("Not set"); + + fireEvent.click(screen.getByTestId("edit-my-limit")); + fireEvent.change(screen.getByTestId("my-limit-input"), { target: { value: "80" } }); + fireEvent.click(screen.getByTestId("save-my-limit")); + + expect(mutate).toHaveBeenCalledWith(80, expect.anything()); + }); + + it("clears a set limit by sending null", () => { + vi.mocked(useMyTeamMember).mockReturnValue({ + data: memberInfo({ effective_budget: 80, self_max_budget: 80, budget_source: "self" }), + isLoading: false, + error: null, + } as ReturnType); + const mutate = mockMutation(); + + renderWithProviders(); + fireEvent.click(screen.getByTestId("clear-my-limit")); + + expect(mutate).toHaveBeenCalledWith(null, expect.anything()); + }); + + it("warns when the typed limit is below the current cycle spend", () => { + vi.mocked(useMyTeamMember).mockReturnValue({ + data: memberInfo({ effective_budget: 100, spend: 50, self_max_budget: null }), + isLoading: false, + error: null, + } as ReturnType); + mockMutation(); + + renderWithProviders(); + fireEvent.click(screen.getByTestId("edit-my-limit")); + fireEvent.change(screen.getByTestId("my-limit-input"), { target: { value: "40" } }); + + expect(screen.getByTestId("below-spend-warning")).toHaveTextContent( + "This is below your current spend of $50.0000. New requests will be blocked until you raise or clear your limit.", + ); + + fireEvent.change(screen.getByTestId("my-limit-input"), { target: { value: "60" } }); + expect(screen.queryByTestId("below-spend-warning")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/MyUserTab.tsx b/ui/litellm-dashboard/src/components/team/MyUserTab.tsx index 9aae6a25dc2..f115eb692f6 100644 --- a/ui/litellm-dashboard/src/components/team/MyUserTab.tsx +++ b/ui/litellm-dashboard/src/components/team/MyUserTab.tsx @@ -2,10 +2,12 @@ import { formatBudgetReset } from "@/utils/budgetUtils"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; import { CircleHelp } from "lucide-react"; -import React from "react"; -import { useMyTeamMember } from "./useMyTeamMember"; +import React, { useState } from "react"; +import { type TeamMemberInfo, useMyTeamMember, useUpdateMySelfBudget } from "./useMyTeamMember"; interface MyUserTabProps { teamId: string; @@ -30,6 +32,110 @@ const formatRateLimit = (value: number | null | undefined): string => { return formatNumberWithCommas(value, 0); }; +const BUDGET_SOURCE_LABELS: Record, string> = { + team_default: "Team default", + custom: "Custom", + self: "Set by you", + none: "None", +}; + +function MyLimitEditor({ + teamId, + selfMaxBudget, + spend, +}: { + teamId: string; + selfMaxBudget: number | null; + spend: number; +}) { + const [editing, setEditing] = useState(false); + const [draft, setDraft] = useState(""); + const mutation = useUpdateMySelfBudget(teamId); + + const parsed = draft.trim() === "" ? NaN : Number(draft); + const draftValid = Number.isFinite(parsed) && parsed >= 0; + const belowSpend = editing && draftValid && parsed < spend; + + const save = () => { + if (!draftValid) return; + mutation.mutate(parsed, { + onSuccess: () => setEditing(false), + }); + }; + + const clear = () => { + mutation.mutate(null, { + onSuccess: () => setEditing(false), + }); + }; + + const startEditing = () => { + setDraft(selfMaxBudget === null ? "" : String(selfMaxBudget)); + mutation.reset(); + setEditing(true); + }; + + if (!editing) { + return ( +
+ + {selfMaxBudget === null ? "Not set" : `$${formatNumber(selfMaxBudget, 4)}`} + + + {selfMaxBudget !== null && ( + + )} +
+ ); + } + + return ( +
+
+ setDraft(e.target.value)} + /> + + +
+ {belowSpend && ( +
+ This is below your current spend of ${formatNumber(spend, 4)}. New requests will be blocked until you raise or + clear your limit. +
+ )} + {mutation.isError && ( +
+ {mutation.error instanceof Error ? mutation.error.message : "Failed to update your limit."} +
+ )} +
+ ); +} + export default function MyUserTab({ teamId }: MyUserTabProps) { const { data, isLoading, error } = useMyTeamMember(teamId); @@ -62,7 +168,9 @@ export default function MyUserTab({ teamId }: MyUserTabProps) { } const budgetTable = data.litellm_budget_table ?? null; - const maxBudget = budgetTable?.max_budget ?? null; + const maxBudget = data.effective_budget ?? null; + const budgetSource = data.budget_source ?? "none"; + const selfMaxBudget = data.self_max_budget ?? null; const spend = data.spend ?? 0; const totalSpend = data.total_spend ?? 0; const tpmLimit = budgetTable?.tpm_limit ?? null; @@ -99,8 +207,13 @@ export default function MyUserTab({ teamId }: MyUserTabProps) { )}

${formatNumber(spend, 4)}

- + of {maxBudget === null ? "Unlimited" : `$${formatNumber(maxBudget, 4)}`} + {budgetSource !== "none" && ( + + {BUDGET_SOURCE_LABELS[budgetSource]} + + )}
{budgetReset &&
Resets {budgetReset}
} @@ -118,6 +231,16 @@ export default function MyUserTab({ teamId }: MyUserTabProps) { + + + {labelWithTooltip( + "My limit", + "A personal limit you set for yourself. It can only lower your team allocation, never raise it.", + )} + + + + {labelWithTooltip("Total Spend (USD)", "Cumulative spend across all budget cycles within this team.")} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 5dfcf1d8e35..60bfbd79345 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -268,6 +268,7 @@ export interface TeamMembership { team_id: string; budget_id: string | null; budget_source: TeamMemberBudgetSource; + self_max_budget?: number | null; spend: number; total_spend: number | null; litellm_budget_table: { diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 1d50a9a4670..62f3ff788ee 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -334,6 +334,45 @@ describe("TeamMembersComponent", () => { expect(screen.getByText("Unlimited")).toBeInTheDocument(); }); + it("shows a read-only Self cap badge when the member has a self_max_budget", () => { + const teamData = createMockTeamData({ + team_memberships: [{ ...createMockTeamData().team_memberships[0], self_max_budget: 250 }], + }); + renderWithProviders( + , + ); + + const badge = screen.getByTestId("member-self-cap"); + expect(badge).toHaveTextContent("Self cap $250.00"); + expect(screen.queryByRole("button", { name: /self cap/i })).not.toBeInTheDocument(); + }); + + it("shows no Self cap badge when the member has no self_max_budget", () => { + renderWithProviders( + , + ); + + expect(screen.queryByTestId("member-self-cap")).not.toBeInTheDocument(); + }); + it("should display No Limits for rate limits when member has no limits", () => { renderWithProviders( { + if (!userId) return null; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + return membership?.self_max_budget ?? null; + }; + const getUserBudget = (userId: string | null): number | null => { if (!userId) return null; const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); @@ -218,9 +224,15 @@ export default function TeamMemberTab({ sortValue: (record: Member) => getUserBudget(record.user_id), render: (record: Member) => { const source = getUserBudgetSource(record.user_id); + const selfCap = getUserSelfCap(record.user_id); return ( + {selfCap !== null && ( + + Self cap ${selfCap.toFixed(2)} + + )} {source !== "none" && ( {BUDGET_SOURCE_LABELS[source]} diff --git a/ui/litellm-dashboard/src/components/team/useMyTeamMember.ts b/ui/litellm-dashboard/src/components/team/useMyTeamMember.ts index 9eec1f135bd..f8369ac024f 100644 --- a/ui/litellm-dashboard/src/components/team/useMyTeamMember.ts +++ b/ui/litellm-dashboard/src/components/team/useMyTeamMember.ts @@ -1,6 +1,7 @@ -import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { useMutation, UseMutationResult, useQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query"; import { deriveErrorMessage, getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { fetchClient } from "@/lib/http/api"; export interface TeamMemberInfo { user_id: string; @@ -9,6 +10,9 @@ export interface TeamMemberInfo { role?: string | null; user_email?: string | null; budget_id?: string | null; + self_max_budget?: number | null; + effective_budget?: number | null; + budget_source?: "team_default" | "custom" | "self" | "none"; spend?: number | null; total_spend?: number | null; litellm_budget_table?: { @@ -63,3 +67,27 @@ export const useMyTeamMember = (teamId: string | null | undefined): UseQueryResu enabled: Boolean(accessToken && teamId), }); }; + +const updateMySelfBudget = async (teamId: string, selfMaxBudget: number | null): Promise => { + const { data, error } = await fetchClient.PATCH("/team/{team_id}/members/me", { + params: { path: { team_id: teamId } }, + body: { self_max_budget: selfMaxBudget }, + }); + + if (error) { + throw new Error(deriveErrorMessage(error)); + } + + return data as TeamMemberInfo; +}; + +export const useUpdateMySelfBudget = (teamId: string): UseMutationResult => { + const queryClient = useQueryClient(); + return useMutation({ + mutationFn: (selfMaxBudget) => updateMySelfBudget(teamId, selfMaxBudget), + onSuccess: (data) => { + queryClient.setQueryData(["team", teamId, "members", "me"], data); + queryClient.invalidateQueries({ queryKey: ["teams"] }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cc0435c3d5e..c3730e3005a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16604,7 +16604,22 @@ export interface paths { delete?: never; options?: never; head?: never; - patch?: never; + /** + * Team Member Me Update Self Budget + * @description Set or clear the caller's own personal spend cap within the team. + * + * The self cap can only lower the member's effective budget below their + * team allocation, never raise it: values above the current admin-set + * allocation are rejected, and the cap binds only while it stays below + * that allocation. Send `self_max_budget: null` to clear it. + * + * Returns 404 if the caller is not a member of the team. + * + * ``` + * curl --location --request PATCH 'http://localhost:4000/team/your_team_id/members/me' --header 'Authorization: Bearer your_api_key_here' --header 'Content-Type: application/json' --data '{"self_max_budget": 80.0}' + * ``` + */ + patch: operations["team_member_me_update_self_budget_team__team_id__members_me_patch"]; trace?: never; }; "/test": { @@ -33563,6 +33578,8 @@ export interface components { budget_id?: string | null; /** Litellm Budget Table */ litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTableFull"] | components["schemas"]["LiteLLM_BudgetTable"] | null; + /** Self Max Budget */ + self_max_budget?: number | null; /** * Spend * @default 0 @@ -43547,10 +43564,20 @@ export interface components { TeamMemberInfoResponse: { /** Budget Id */ budget_id?: string | null; + /** + * Budget Source + * @default none + * @enum {string} + */ + budget_source: "team_default" | "custom" | "self" | "none"; + /** Effective Budget */ + effective_budget?: number | null; /** Litellm Budget Table */ litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTableFull"] | components["schemas"]["LiteLLM_BudgetTable"] | null; /** Role */ role?: string | null; + /** Self Max Budget */ + self_max_budget?: number | null; /** * Spend * @default 0 @@ -43596,6 +43623,16 @@ export interface components { /** User Id */ user_id: string; }; + /** + * TeamMemberSelfBudgetUpdateRequest + * @description Body of `PATCH /team/{team_id}/members/me`. `self_max_budget` is + * required: a non-negative number sets the caller's personal cap, null + * clears it. + */ + TeamMemberSelfBudgetUpdateRequest: { + /** Self Max Budget */ + self_max_budget: number | null; + }; /** TeamMemberUpdateRequest */ TeamMemberUpdateRequest: { /** @@ -68135,6 +68172,41 @@ export interface operations { }; }; }; + team_member_me_update_self_budget_team__team_id__members_me_patch: { + parameters: { + query?: never; + header?: never; + path: { + team_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["TeamMemberSelfBudgetUpdateRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["TeamMemberInfoResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; test_endpoint_test_get: { parameters: { query?: never;