scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(proxy): let team members set their own spend cap within their team allocation

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mrinal 2026-09-30 09:22:47 +00:00
parent 79756cbb9b
commit 6339f9fef8
20 changed files with 1696 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260929000000_add_team_membership_self_max_budget/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN IF NOT EXISTS "self_max_budget" DOUBLE PRECISION;

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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])
}

11
litellm/models/team_membership.py
View file

@ -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

102
litellm/proxy/auth/auth_checks.py
View file

@ -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,

284
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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,
)

1
litellm/proxy/schema.prisma
View file

@ -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])
}

10
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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

17
litellm/types/proxy/management_endpoints/team_endpoints.py
View file

@ -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):

1
schema.prisma
View file

@ -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])
}

1
terraform/provider/tools/endpointaudit/coverage_allowlist.txt
View file

@ -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,

243
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -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

640
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -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)

84
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

@ -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

92
ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx
View file

@ -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<string, unknown> = {}) => ({
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<typeof useUpdateMySelfBudget>);
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<typeof useMyTeamMember>);
mockMutation();
renderWithProviders(<MyUserTab teamId="team-1" />);
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<typeof useMyTeamMember>);
const mutate = mockMutation();
renderWithProviders(<MyUserTab teamId="team-1" />);
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<typeof useMyTeamMember>);
const mutate = mockMutation();
renderWithProviders(<MyUserTab teamId="team-1" />);
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<typeof useMyTeamMember>);
mockMutation();
renderWithProviders(<MyUserTab teamId="team-1" />);
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();
});
});

131
ui/litellm-dashboard/src/components/team/MyUserTab.tsx
View file

@ -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<NonNullable<TeamMemberInfo["budget_source"]>, 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 (
<div className="mt-2 flex items-center gap-2">
<span className="text-xl font-semibold" data-testid="my-limit-value">
{selfMaxBudget === null ? "Not set" : `$${formatNumber(selfMaxBudget, 4)}`}
</span>
<Button variant="outline" size="xs" data-testid="edit-my-limit" onClick={startEditing}>
Edit
</Button>
{selfMaxBudget !== null && (
<Button variant="link" size="xs" data-testid="clear-my-limit" disabled={mutation.isPending} onClick={clear}>
Clear
</Button>
)}
</div>
);
}
return (
<div className="mt-2">
<div className="flex items-center gap-2">
<Input
type="number"
min={0}
step="any"
value={draft}
data-testid="my-limit-input"
aria-label="My limit"
className="w-40"
onChange={(e) => setDraft(e.target.value)}
/>
<Button size="xs" data-testid="save-my-limit" disabled={!draftValid || mutation.isPending} onClick={save}>
Save
</Button>
<Button
variant="outline"
size="xs"
data-testid="cancel-my-limit"
onClick={() => {
mutation.reset();
setEditing(false);
}}
>
Cancel
</Button>
</div>
{belowSpend && (
<div className="mt-1 text-amber-600" data-testid="below-spend-warning">
This is below your current spend of ${formatNumber(spend, 4)}. New requests will be blocked until you raise or
clear your limit.
</div>
)}
{mutation.isError && (
<div className="mt-1 text-destructive" data-testid="my-limit-error">
{mutation.error instanceof Error ? mutation.error.message : "Failed to update your limit."}
</div>
)}
</div>
);
}
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) {
)}
<div className="mt-2">
<h3 className="text-2xl font-semibold">${formatNumber(spend, 4)}</h3>
<span className="text-muted-foreground">
<span className="inline-flex items-center gap-2 text-muted-foreground">
of {maxBudget === null ? "Unlimited" : `$${formatNumber(maxBudget, 4)}`}
{budgetSource !== "none" && (
<Badge variant={budgetSource === "self" ? "outline" : "secondary"} data-testid="budget-source-badge">
{BUDGET_SOURCE_LABELS[budgetSource]}
</Badge>
)}
</span>
</div>
{budgetReset && <div className="mt-1 text-muted-foreground">Resets {budgetReset}</div>}
@ -118,6 +231,16 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
</CardContent>
</Card>
<Card>
<CardContent>
{labelWithTooltip(
"My limit",
"A personal limit you set for yourself. It can only lower your team allocation, never raise it.",
)}
<MyLimitEditor teamId={teamId} selfMaxBudget={selfMaxBudget} spend={spend} />
</CardContent>
</Card>
<Card>
<CardContent>
{labelWithTooltip("Total Spend (USD)", "Cumulative spend across all budget cycles within this team.")}

1
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -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: {

39
ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx
View file

@ -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(
<TeamMembersComponent
teamData={teamData}
canEditTeam={false}
handleMemberDelete={mockHandleMemberDelete}
onMemberSpendReset={mockOnMemberSpendReset}
onMemberBudgetReset={mockOnMemberBudgetReset}
setSelectedEditMember={mockSetSelectedEditMember}
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
/>,
);
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(
<TeamMembersComponent
teamData={createMockTeamData()}
canEditTeam={false}
handleMemberDelete={mockHandleMemberDelete}
onMemberSpendReset={mockOnMemberSpendReset}
onMemberBudgetReset={mockOnMemberBudgetReset}
setSelectedEditMember={mockSetSelectedEditMember}
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
/>,
);
expect(screen.queryByTestId("member-self-cap")).not.toBeInTheDocument();
});
it("should display No Limits for rate limits when member has no limits", () => {
renderWithProviders(
<TeamMembersComponent

12
ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx
View file

@ -103,6 +103,12 @@ export default function TeamMemberTab({
return membership?.budget_source ?? "none";
};
const getUserSelfCap = (userId: string | null): number | null => {
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 (
<span className="flex items-center justify-end gap-2">
<MoneyCell value={getUserBudget(record.user_id)} decimals={2} emptyText="Unlimited" showZero />
{selfCap !== null && (
<Badge variant="outline" data-testid="member-self-cap">
Self cap ${selfCap.toFixed(2)}
</Badge>
)}
{source !== "none" && (
<Badge variant={source === "custom" ? "outline" : "secondary"} data-testid="member-budget-source">
{BUDGET_SOURCE_LABELS[source]}

30
ui/litellm-dashboard/src/components/team/useMyTeamMember.ts
View file

@ -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<TeamMemberInfo> => {
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<TeamMemberInfo, Error, number | null> => {
const queryClient = useQueryClient();
return useMutation<TeamMemberInfo, Error, number | null>({
mutationFn: (selfMaxBudget) => updateMySelfBudget(teamId, selfMaxBudget),
onSuccess: (data) => {
queryClient.setQueryData(["team", teamId, "members", "me"], data);
queryClient.invalidateQueries({ queryKey: ["teams"] });
},
});
};

74
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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;