ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ui): add budget duration to edit team member form (#29717)

Browse Source 
* feat(ui): add budget duration to edit team member form

Editing a team member created a member budget with no duration, so the
budget never reset. This threads a budget reset period through the edit
flow end to end and reuses the shared duration dropdown so the options
stay in sync with the rest of the UI.

Resolves LIT-2651

* fix(proxy): validate member budget_duration and persist clears

Reject budget_duration values that can't be parsed, are non-positive, or overflow date math before any write, so a bad value can't be persisted and later crash the budget reset job.

Clearing the budget duration in the edit-member form now sends null and clears the column end to end, so the dropdown's clear control reflects a real change instead of being a no-op

* chore(ui): regenerate schema.d.ts for member budget_duration

Adds budget_duration to TeamMemberUpdateRequest/Response in the generated dashboard types so the Check UI API Types Sync gate passes
This commit is contained in:
ryan-crabbe-berri 2026-06-06 17:24:55 -07:00 committed by GitHub
parent aeb55e7a11
commit f31d059aa3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 512 additions and 315 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
litellm/proxy/_types.py
View File

@ -4125,6 +4125,10 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
rpm_limit: Optional[int] = Field(
default=None, description="Requests per minute limit for this team member"
)
budget_duration: Optional[str] = Field(
default=None,
description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.",
)
allowed_models: Optional[List[str]] = Field(
default=None,
description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.",
@ -4136,6 +4140,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse):
max_budget_in_team: Optional[float] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
budget_duration: Optional[str] = None
allowed_models: Optional[List[str]] = None

161
litellm/proxy/management_endpoints/common_utils.py
View File

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from fastapi import HTTPException, status
from pydantic import BaseModel
@ -19,6 +19,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
user_api_key_has_admin_view as _user_has_admin_view, # noqa: F401 re-exported
)
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.utils import _premium_user_check
if TYPE_CHECKING:
@ -400,121 +401,127 @@ def _set_object_metadata_field(
object_data.metadata[field_name] = value
_TEAM_MEMBER_BUDGET_LIMIT_FIELDS = (
"max_budget",
"soft_budget",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
"model_max_budget",
"budget_duration",
"allowed_models",
)
def _is_set_budget_value(value: Any) -> bool:
if value is None:
return False
if isinstance(value, list) and len(value) == 0:
return False
return True
def _has_meaningful_budget_limit(budget_values: Dict[str, Any]) -> bool:
"""A budget is meaningful if at least one limit is actually set; an empty
list (no model restriction) and None both count as unset."""
return any(
_is_set_budget_value(budget_values.get(field))
for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS
)
async def _upsert_budget_and_membership(
tx,
*,
team_id: str,
user_id: str,
max_budget: Optional[float],
existing_budget_id: Optional[str],
user_api_key_dict: UserAPIKeyAuth,
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
allowed_models: Optional[List[str]] = None,
budget_patch: Dict[str, Any],
team_default_budget_id: Optional[str] = None,
):
"""
Helper function to Create/Update or Delete the budget within the team membership
Args:
tx: The transaction object
team_id: The ID of the team
user_id: The ID of the user
max_budget: The maximum budget for the team
existing_budget_id: The ID of the existing budget, if any
user_api_key_dict: User API Key dictionary containing user information
tpm_limit: Tokens per minute limit for the team member
rpm_limit: Requests per minute limit for the team member
allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce.
team_default_budget_id: The team's shared default member budget id (from
team metadata.team_member_budget_id), if any. When the membership's
existing_budget_id matches this, we clone-on-write so editing one
member's budget does not mutate the shared default (and therefore
every other member who still points at it).
Apply a merge-patch of per-member budget fields to a team membership.
If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership.
If any of these values exist, a budget is updated or created and linked to the team membership.
``budget_patch`` holds only the budget columns the caller explicitly sent
(RFC 7396 semantics): a value sets the column, ``None`` clears it, and a
column that is absent from the dict is left untouched. Once the patch is
applied, if the budget has no meaningful limit left the member's private
budget is disconnected so they fall back to the team default.
``team_default_budget_id`` is the team's shared default member budget id
(from team metadata.team_member_budget_id). When the membership still
points at it, we clone-on-write so editing one member's budget does not
mutate the shared default that every other member points at.
"""
if (
max_budget is None
and tpm_limit is None
and rpm_limit is None
and allowed_models is None
):
# disconnect the budget since all limits are None
await tx.litellm_teammembership.update(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
data={"litellm_budget_table": {"disconnect": True}},
)
if not budget_patch:
return
write_data = dict(budget_patch)
if "budget_duration" in write_data:
duration = write_data["budget_duration"]
write_data["budget_reset_at"] = (
get_budget_reset_time(budget_duration=duration)
if duration is not None
else None
)
is_shared_default = (
existing_budget_id is not None
and team_default_budget_id is not None
and existing_budget_id == team_default_budget_id
)
async def _disconnect():
await tx.litellm_teammembership.update(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
data={"litellm_budget_table": {"disconnect": True}},
)
if existing_budget_id is not None and not is_shared_default:
# Update the existing budget in-place to preserve fields not being changed.
# Only write fields that the caller explicitly provided (non-None).
update_data: Dict[str, Any] = {
"updated_by": user_api_key_dict.user_id or "",
}
if max_budget is not None:
update_data["max_budget"] = max_budget
if tpm_limit is not None:
update_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
update_data["rpm_limit"] = rpm_limit
if allowed_models is not None:
update_data["allowed_models"] = allowed_models
existing_budget = await tx.litellm_budgettable.find_unique(
where={"budget_id": existing_budget_id}
)
merged = existing_budget.model_dump() if existing_budget is not None else {}
merged.update(write_data)
if not _has_meaningful_budget_limit(merged):
await _disconnect()
return
await tx.litellm_budgettable.update(
where={"budget_id": existing_budget_id},
data=update_data,
data={"updated_by": user_api_key_dict.user_id or "", **write_data},
)
return
# Either there is no existing budget, OR the membership is still pointing
# at the team's shared default member budget. In both cases we create a
# NEW private budget for this user and (re)link the membership to it.
create_data: Dict[str, Any] = {
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
}
# If we're forking off the shared default, seed the new row with the
# default's values so fields the caller did not change carry over.
if is_shared_default:
default_budget_row = await tx.litellm_budgettable.find_unique(
where={"budget_id": existing_budget_id}
)
if default_budget_row is not None:
default_budget_dict = default_budget_row.model_dump()
for field in (
"max_budget",
"soft_budget",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
"model_max_budget",
"budget_duration",
"allowed_models",
):
for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS:
value = default_budget_dict.get(field)
if value is None:
continue
if isinstance(value, list) and len(value) == 0:
continue
create_data[field] = value
if _is_set_budget_value(value):
create_data[field] = value
# Caller-provided values take precedence over the cloned defaults.
if max_budget is not None:
create_data["max_budget"] = max_budget
if tpm_limit is not None:
create_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
create_data["rpm_limit"] = rpm_limit
if allowed_models is not None:
create_data["allowed_models"] = allowed_models
create_data.update(write_data)
if create_data.get("budget_duration") is not None:
create_data["budget_reset_at"] = get_budget_reset_time(
budget_duration=create_data["budget_duration"]
)
else:
create_data.pop("budget_reset_at", None)
if not _has_meaningful_budget_limit(create_data):
if existing_budget_id is not None:
await _disconnect()
return
new_budget = await tx.litellm_budgettable.create(
data=create_data,

55
litellm/proxy/management_endpoints/team_endpoints.py
View File

@ -2733,6 +2733,52 @@ async def team_member_delete(
return existing_team_row
_MEMBER_BUDGET_PATCH_FIELDS = {
"max_budget_in_team": "max_budget",
"tpm_limit": "tpm_limit",
"rpm_limit": "rpm_limit",
"budget_duration": "budget_duration",
"allowed_models": "allowed_models",
}
def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]:
"""Map the budget fields the request actually set (merge-patch: a sent
value updates, an explicit null clears, an absent field is left untouched)
to their budget-table columns."""
provided = data.model_dump(exclude_unset=True)
return {
column: provided[request_field]
for request_field, column in _MEMBER_BUDGET_PATCH_FIELDS.items()
if request_field in provided
}
def _validate_budget_duration(budget_duration: Optional[str]) -> None:
"""Reject budget durations that can't be parsed, are non-positive, or
overflow date math, so a bad value can't be persisted and later crash the
budget reset job."""
if budget_duration is None:
return
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
try:
if duration_in_seconds(budget_duration) <= 0:
raise ValueError("budget_duration must be positive")
get_budget_reset_time(budget_duration=budget_duration)
except (ValueError, OverflowError):
raise HTTPException(
status_code=400,
detail={
"error": "Invalid budget_duration '{}'. Use a format like '1h', '24h', '7d', or '30d'.".format(
budget_duration
)
},
)
@router.post(
"/team/member_update",
tags=["team management"],
@ -2770,6 +2816,8 @@ async def team_member_update(
detail={"error": "Either user_id or user_email needs to be passed in"},
)
_validate_budget_duration(data.budget_duration)
_existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": data.team_id}
)
@ -2843,17 +2891,15 @@ async def team_member_update(
team_default_budget_id = raw_default_budget_id
### upsert new budget
budget_patch = _build_member_budget_patch(data)
async with prisma_client.db.tx() as tx:
await _upsert_budget_and_membership(
tx=tx,
team_id=data.team_id,
user_id=received_user_id,
max_budget=data.max_budget_in_team,
existing_budget_id=identified_budget_id,
user_api_key_dict=user_api_key_dict,
tpm_limit=data.tpm_limit,
rpm_limit=data.rpm_limit,
allowed_models=data.allowed_models,
budget_patch=budget_patch,
team_default_budget_id=team_default_budget_id,
)
@ -2887,6 +2933,7 @@ async def team_member_update(
max_budget_in_team=data.max_budget_in_team,
tpm_limit=data.tpm_limit,
rpm_limit=data.rpm_limit,
budget_duration=data.budget_duration,
allowed_models=data.allowed_models,
)

417
tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py
View File

@ -1,5 +1,6 @@
# tests/litellm/proxy/common_utils/test_upsert_budget_membership.py
import types
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -19,15 +20,13 @@ def mock_tx():
Builds an object that looks just enough like the Prisma tx you use
inside _upsert_budget_and_membership.
"""
# membership “table”
membership = MagicMock()
membership.update = AsyncMock()
membership.upsert = AsyncMock()
# budget “table”
budget = MagicMock()
budget.update = AsyncMock()
# budget.create returns a fake row that has .budget_id
budget.find_unique = AsyncMock(return_value=None)
budget.create = AsyncMock(
return_value=types.SimpleNamespace(budget_id="new-budget-123")
)
@ -44,16 +43,57 @@ def fake_user():
return types.SimpleNamespace(user_id="tester@example.com")
# TEST: max_budget is None, disconnect only
def budget_row(**fields):
"""A fake litellm_budgettable row whose model_dump returns the given fields."""
row = MagicMock()
row.model_dump.return_value = fields
return row
def assert_future_reset_time(value):
"""A budget_reset_at must be a timezone-aware datetime in the future, so the
member's budget actually rolls over and the UI shows a reset date instead of
waiting for the reset cron to backfill it."""
assert isinstance(value, datetime)
assert value.tzinfo is not None
assert value > datetime.now(timezone.utc)
# TEST: an empty patch (caller sent no budget fields) leaves everything alone.
# This is the merge-patch contract: absent != clear. Updating only a member's
# role must not silently wipe their budget.
@pytest.mark.asyncio
async def test_upsert_disconnect(mock_tx, fake_user):
async def test_empty_patch_is_noop(mock_tx, fake_user):
await _upsert_budget_and_membership(
mock_tx,
team_id="team-1",
user_id="user-1",
max_budget=None,
existing_budget_id=None,
existing_budget_id="bud-1",
user_api_key_dict=fake_user,
budget_patch={},
)
mock_tx.litellm_teammembership.update.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_budgettable.create.assert_not_called()
# TEST: clearing every limit on a member's private budget disconnects it, so the
# member falls back to the team default instead of keeping an empty private row.
@pytest.mark.asyncio
async def test_clearing_all_limits_disconnects(mock_tx, fake_user):
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(max_budget=100.0)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-1",
user_id="user-1",
existing_budget_id="bud-1",
user_api_key_dict=fake_user,
budget_patch={"max_budget": None},
)
mock_tx.litellm_teammembership.update.assert_awaited_once_with(
@ -62,205 +102,114 @@ async def test_upsert_disconnect(mock_tx, fake_user):
)
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: existing budget id → updates budget in-place (current behavior)
# TEST: clearing one field on a budget that still has another limit updates in
# place (clears just that column + its reset time) and does NOT disconnect.
@pytest.mark.asyncio
async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user):
"""
Test that when existing_budget_id is provided, the function updates the budget in-place.
"""
await _upsert_budget_and_membership(
mock_tx,
team_id="team-2",
user_id="user-2",
max_budget=42.0,
existing_budget_id="bud-999",
user_api_key_dict=fake_user,
async def test_clear_one_field_keeps_others(mock_tx, fake_user):
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(max_budget=100.0, budget_duration="24h")
)
# Should update the existing budget, not create a new one
await _upsert_budget_and_membership(
mock_tx,
team_id="team-1",
user_id="user-1",
existing_budget_id="bud-1",
user_api_key_dict=fake_user,
budget_patch={"budget_duration": None},
)
mock_tx.litellm_teammembership.update.assert_not_called()
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "bud-999"},
where={"budget_id": "bud-1"},
data={
"max_budget": 42.0,
"updated_by": fake_user.user_id,
"budget_duration": None,
"budget_reset_at": None,
},
)
# Should NOT create a new budget or touch membership
# TEST: setting budget_duration in place writes the duration AND a future
# budget_reset_at, so the budget rolls over without waiting for the reset cron.
@pytest.mark.asyncio
async def test_update_in_place_seeds_reset_at(mock_tx, fake_user):
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(max_budget=20.0)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-dur",
user_id="user-dur",
existing_budget_id="bud-dur",
user_api_key_dict=fake_user,
budget_patch={"budget_duration": "30d"},
)
mock_tx.litellm_budgettable.update.assert_awaited_once()
call = mock_tx.litellm_budgettable.update.await_args
assert call.kwargs["where"] == {"budget_id": "bud-dur"}
data = call.kwargs["data"]
assert data["budget_duration"] == "30d"
assert data["updated_by"] == fake_user.user_id
assert_future_reset_time(data["budget_reset_at"])
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
mock_tx.litellm_teammembership.update.assert_not_called()
# TEST: create new budget and link membership
# TEST: updating a single limit in place only writes that field; an untouched
# budget_duration must not get a (re)computed reset time.
@pytest.mark.asyncio
async def test_upsert_create_and_link(mock_tx, fake_user):
async def test_update_in_place_single_field_leaves_reset_at_alone(mock_tx, fake_user):
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(max_budget=50.0)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-3",
user_id="user-3",
max_budget=99.9,
existing_budget_id=None,
team_id="team-rpm",
user_id="user-rpm",
existing_budget_id="bud-rpm",
user_api_key_dict=fake_user,
budget_patch={"rpm_limit": 100},
)
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
data={
"max_budget": 99.9,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "bud-rpm"},
data={"updated_by": fake_user.user_id, "rpm_limit": 100},
)
# Budget ID returned by the mocked create()
bid = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-3", "team_id": "team-3"}},
data={
"create": {
"user_id": "user-3",
"team_id": "team-3",
"litellm_budget_table": {"connect": {"budget_id": bid}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": bid}},
},
},
)
mock_tx.litellm_teammembership.update.assert_not_called()
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_budgettable.create.assert_not_called()
# TEST: create new budget and link membership, then create another new budget
# TEST: with no existing budget, a duration-only patch creates a budget carrying
# the duration and a future reset time, then links the membership.
@pytest.mark.asyncio
async def test_upsert_create_then_create_another(mock_tx, fake_user):
"""
Test that multiple calls to _upsert_budget_and_membership create separate budgets,
reflecting the current implementation behavior.
"""
# FIRST CALL – create new budget and link membership
async def test_create_seeds_reset_at_and_links(mock_tx, fake_user):
await _upsert_budget_and_membership(
mock_tx,
team_id="team-42",
user_id="user-42",
max_budget=10.0,
team_id="team-new",
user_id="user-new",
existing_budget_id=None,
user_api_key_dict=fake_user,
budget_patch={"budget_duration": "7d"},
)
# capture the budget id that create() returned
created_bid = mock_tx.litellm_budgettable.create.return_value.budget_id
# sanity: we really did the create + upsert path
mock_tx.litellm_budgettable.create.assert_awaited_once()
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert data["budget_duration"] == "7d"
assert data["created_by"] == fake_user.user_id
assert data["updated_by"] == fake_user.user_id
assert_future_reset_time(data["budget_reset_at"])
# SECOND CALL – reset call history; this time we supply the existing budget_id
mock_tx.litellm_budgettable.create.reset_mock()
mock_tx.litellm_teammembership.upsert.reset_mock()
mock_tx.litellm_budgettable.update.reset_mock()
await _upsert_budget_and_membership(
mock_tx,
team_id="team-42",
user_id="user-42",
max_budget=25.0,
existing_budget_id=created_bid, # now used: triggers in-place update
user_api_key_dict=fake_user,
)
# Should update the existing budget in-place, not create a new one
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": created_bid},
data={
"max_budget": 25.0,
"updated_by": fake_user.user_id,
},
)
# Should NOT create a new budget or touch membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: update rpm_limit for member with existing budget_id → updates in-place
@pytest.mark.asyncio
async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user):
"""
Test that updating rpm_limit for a member with an existing budget_id
updates the existing budget in-place (not creates a new one).
"""
existing_budget_id = "existing-budget-456"
await _upsert_budget_and_membership(
mock_tx,
team_id="team-rpm-test",
user_id="user-rpm-test",
max_budget=50.0,
existing_budget_id=existing_budget_id,
user_api_key_dict=fake_user,
tpm_limit=1000,
rpm_limit=100,
)
# Should update the existing budget with all specified limits
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": existing_budget_id},
data={
"max_budget": 50.0,
"tpm_limit": 1000,
"rpm_limit": 100,
"updated_by": fake_user.user_id,
},
)
# Should NOT create a new budget or touch membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: create new budget with only rpm_limit (no max_budget)
@pytest.mark.asyncio
async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user):
"""
Test that setting only rpm_limit creates a new budget with just the rpm_limit.
"""
await _upsert_budget_and_membership(
mock_tx,
team_id="team-rpm-only",
user_id="user-rpm-only",
max_budget=None,
existing_budget_id=None,
user_api_key_dict=fake_user,
rpm_limit=50,
)
# Should create a new budget with only rpm_limit
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
data={
"rpm_limit": 50,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={
"user_id_team_id": {"user_id": "user-rpm-only", "team_id": "team-rpm-only"}
},
where={"user_id_team_id": {"user_id": "user-new", "team_id": "team-new"}},
data={
"create": {
"user_id": "user-rpm-only",
"team_id": "team-rpm-only",
"user_id": "user-new",
"team_id": "team-new",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
@ -270,60 +219,48 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user):
)
# TEST: clone-on-write when membership still points at the team's shared default budget
# TEST: clone-on-write when the membership still points at the team's shared
# default budget. Editing this member must fork a private budget instead of
# mutating the shared row, and cloning a duration must seed a fresh reset time.
@pytest.mark.asyncio
async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user):
"""
When a member's existing budget_id is the same row as the team's shared
default member budget, updating that member's budget must NOT mutate the
shared row. Instead we should create a new private budget for this member
(seeded with the default's values) and re-link the membership to it.
"""
async def test_clone_on_write_from_shared_default(mock_tx, fake_user):
shared_default_id = "team-default-budget-1"
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(
budget_id=shared_default_id,
max_budget=200.0,
soft_budget=None,
max_parallel_requests=None,
tpm_limit=500,
rpm_limit=None,
model_max_budget=None,
budget_duration="1d",
allowed_models=[],
)
)
# Default budget row in the DB: $200 cap, daily reset, 500 tpm.
default_row = MagicMock()
default_row.model_dump.return_value = {
"budget_id": shared_default_id,
"max_budget": 200.0,
"soft_budget": None,
"max_parallel_requests": None,
"tpm_limit": 500,
"rpm_limit": None,
"model_max_budget": None,
"budget_duration": "1d",
"allowed_models": [],
}
mock_tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row)
# Caller is changing only this member's max_budget.
await _upsert_budget_and_membership(
mock_tx,
team_id="team-shared",
user_id="user-shared",
max_budget=50.0,
existing_budget_id=shared_default_id,
user_api_key_dict=fake_user,
budget_patch={"max_budget": 50.0},
team_default_budget_id=shared_default_id,
)
# Must NOT touch the shared default row in place.
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_budgettable.create.assert_awaited_once()
create_data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert_future_reset_time(create_data.pop("budget_reset_at"))
assert create_data == {
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"max_budget": 50.0, # caller wins
"tpm_limit": 500, # cloned from default
"budget_duration": "1d", # cloned from default
}
# Must create a new private budget seeded with the default's values,
# with the caller's max_budget overriding the cloned default.
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
data={
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"max_budget": 50.0, # caller wins
"tpm_limit": 500, # cloned from default
"budget_duration": "1d", # cloned from default
},
include={"team_membership": True},
)
# Membership must be re-linked to the new private budget.
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-shared", "team_id": "team-shared"}},
@ -340,32 +277,64 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user)
)
# TEST: when team default exists but member already has their own budget, in-place update
# TEST: forking the shared default while clearing its duration must drop the
# duration (and not carry a reset time) on the new private budget.
@pytest.mark.asyncio
async def test_upsert_updates_in_place_when_member_has_private_budget(
mock_tx, fake_user
):
"""
If the member's budget_id is different from the team's shared default
(i.e. they already have a private budget), we should keep the current
in-place behavior and not allocate a new row.
"""
async def test_clone_on_write_clears_duration(mock_tx, fake_user):
shared_default_id = "team-default-budget-1"
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(
budget_id=shared_default_id,
max_budget=200.0,
tpm_limit=500,
budget_duration="1d",
allowed_models=[],
)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-shared",
user_id="user-shared",
existing_budget_id=shared_default_id,
user_api_key_dict=fake_user,
budget_patch={"budget_duration": None},
team_default_budget_id=shared_default_id,
)
mock_tx.litellm_budgettable.update.assert_not_called()
create_data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert create_data == {
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"max_budget": 200.0,
"tpm_limit": 500,
"budget_duration": None,
}
assert "budget_reset_at" not in create_data
# TEST: when the member already has their own private budget (different from the
# team default), we update it in place rather than forking another row.
@pytest.mark.asyncio
async def test_private_budget_updates_in_place(mock_tx, fake_user):
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(max_budget=10.0)
)
await _upsert_budget_and_membership(
mock_tx,
team_id="team-mixed",
user_id="user-private",
max_budget=75.0,
existing_budget_id="private-budget-xyz",
user_api_key_dict=fake_user,
budget_patch={"max_budget": 75.0},
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "private-budget-xyz"},
data={
"max_budget": 75.0,
"updated_by": fake_user.user_id,
},
data={"max_budget": 75.0, "updated_by": fake_user.user_id},
)
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()

142
tests/test_litellm/proxy/test_team_member_update.py
View File

@ -1,9 +1,19 @@
import types
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from starlette.requests import Request
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import TeamMemberUpdateRequest
import litellm.proxy.management_endpoints.team_endpoints as team_endpoints
from litellm.proxy._types import (
LiteLLM_TeamTable,
LitellmUserRoles,
Member,
TeamMemberUpdateRequest,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.team_endpoints import team_member_update
@ -38,3 +48,133 @@ async def test_ateam_member_update_admin_requires_premium(monkeypatch):
"Pricing: https://www.litellm.ai/#pricing"
)
assert exc_info.value.detail == expected_msg
@pytest.fixture
def happy_path_upsert(monkeypatch):
"""Stub out the DB and the budget upsert so a team_member_update call reaches
_upsert_budget_and_membership, and hand back that mock to inspect the patch."""
team_row = LiteLLM_TeamTable(
team_id="team-1234",
members_with_roles=[Member(user_id="user-1", role="user")],
metadata={},
)
prisma_client = MagicMock()
prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
prisma_client.db.litellm_teamtable.update = AsyncMock()
class _FakeTx:
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return False
prisma_client.db.tx = MagicMock(return_value=_FakeTx())
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "premium_user", False)
monkeypatch.setattr(
team_endpoints,
"team_info",
AsyncMock(
return_value={
"team_info": team_row,
"team_memberships": [
types.SimpleNamespace(user_id="user-1", budget_id="bud-1")
],
}
),
)
upsert_mock = AsyncMock()
monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock)
return upsert_mock
def _member_update_request(**overrides):
data = TeamMemberUpdateRequest(
team_id="team-1234", user_id="user-1", role="user", **overrides
)
request = Request({"type": "http", "method": "POST", "path": "/team/member_update"})
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id="admin")
return data, request, auth
@pytest.mark.asyncio
async def test_team_member_update_sends_provided_fields_as_patch(happy_path_upsert):
"""Fields the request sets must reach _upsert_budget_and_membership as a
budget patch, otherwise the member budget is never written/reset."""
data, request, auth = _member_update_request(
max_budget_in_team=10.0, budget_duration="30d"
)
response = await team_member_update(data, request, auth)
happy_path_upsert.assert_awaited_once()
assert happy_path_upsert.await_args.kwargs["budget_patch"] == {
"max_budget": 10.0,
"budget_duration": "30d",
}
assert response.budget_duration == "30d"
@pytest.mark.asyncio
async def test_team_member_update_explicit_null_clears_field(happy_path_upsert):
"""An explicitly-null field must be forwarded as None so the column is
cleared, rather than silently dropped."""
data, request, auth = _member_update_request(budget_duration=None)
await team_member_update(data, request, auth)
assert happy_path_upsert.await_args.kwargs["budget_patch"] == {
"budget_duration": None
}
@pytest.mark.asyncio
async def test_team_member_update_omits_unset_fields_from_patch(happy_path_upsert):
"""A request that touches no budget fields must produce an empty patch so the
member's existing budget is left untouched."""
data, request, auth = _member_update_request()
await team_member_update(data, request, auth)
assert happy_path_upsert.await_args.kwargs["budget_patch"] == {}
@pytest.mark.parametrize(
"bad_duration",
[
"not-a-duration", # unparseable garbage
"10x", # unsupported unit
"0d", # zero-length window
"999999999999999999999999d", # overflows datetime math
],
)
@pytest.mark.asyncio
async def test_team_member_update_rejects_invalid_budget_duration(
monkeypatch, bad_duration
):
"""An invalid budget_duration must be rejected with a 400 before any DB
write, so it can never be persisted and later break the budget reset job."""
monkeypatch.setattr(proxy_server, "prisma_client", object())
monkeypatch.setattr(proxy_server, "premium_user", False)
upsert_mock = AsyncMock()
monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock)
data = TeamMemberUpdateRequest(
team_id="team-1234",
user_id="user-1",
role="user",
budget_duration=bad_duration,
)
request = Request({"type": "http", "method": "POST", "path": "/team/member_update"})
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id="admin")
with pytest.raises(HTTPException) as exc_info:
await team_member_update(data, request, auth)
assert exc_info.value.status_code == 400
assert "budget_duration" in str(exc_info.value.detail)
upsert_mock.assert_not_called()

18
ui/litellm-dashboard/src/components/networking.tsx
View File

@ -2822,6 +2822,7 @@ export interface Member {
max_budget_in_team?: number | null;
tpm_limit?: number | null;
rpm_limit?: number | null;
budget_duration?: string | null;
allowed_models?: string[] | null;
}
@ -2949,18 +2950,21 @@ export const teamMemberUpdateCall = async (
user_id: formValues.user_id,
};
// Add optional budget and rate limit fields
const orNull = (value: unknown) => (value === undefined || value === null || value === "" ? null : value);
if (formValues.user_email !== undefined) {
requestBody.user_email = formValues.user_email;
}
if (formValues.max_budget_in_team !== undefined && formValues.max_budget_in_team !== null) {
requestBody.max_budget_in_team = formValues.max_budget_in_team;
if ("max_budget_in_team" in formValues) {
requestBody.max_budget_in_team = orNull(formValues.max_budget_in_team);
}
if (formValues.tpm_limit !== undefined && formValues.tpm_limit !== null) {
requestBody.tpm_limit = formValues.tpm_limit;
if ("tpm_limit" in formValues) {
requestBody.tpm_limit = orNull(formValues.tpm_limit);
}
if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) {
requestBody.rpm_limit = formValues.rpm_limit;
if ("rpm_limit" in formValues) {
requestBody.rpm_limit = orNull(formValues.rpm_limit);
}
if ("budget_duration" in formValues) {
requestBody.budget_duration = orNull(formValues.budget_duration);
}
if (formValues.allowed_models !== undefined) {
requestBody.allowed_models = formValues.allowed_models;

8
ui/litellm-dashboard/src/components/team/EditMembership.tsx
View File

@ -2,6 +2,7 @@ import { Text, TextInput } from "@tremor/react";
import { Button as AntButton, Form, Modal, Select } from "antd";
import React, { useEffect, useState } from "react";
import NumericalInput from "../shared/numerical_input";
import BudgetDurationDropdown from "../common_components/budget_duration_dropdown";
interface BaseMember {
user_email?: string;
@ -21,7 +22,7 @@ interface ModalConfig {
additionalFields?: Array<{
name: string;
label: string | React.ReactNode;
type: "input" | "select" | "numerical" | "multi-select";
type: "input" | "select" | "numerical" | "multi-select" | "budget-duration";
options?: Array<{ label: string; value: string }>;
rules?: any[];
step?: number;
@ -65,6 +66,7 @@ const MemberModal = <T extends BaseMember>({
max_budget_in_team: (initialData as any).max_budget_in_team || null,
tpm_limit: (initialData as any).tpm_limit || null,
rpm_limit: (initialData as any).rpm_limit || null,
budget_duration: (initialData as any).budget_duration || null,
// Keep array values for multi-select fields
allowed_models: (initialData as any).allowed_models || [],
};
@ -117,7 +119,7 @@ const MemberModal = <T extends BaseMember>({
const renderField = (field: {
name: string;
label: string | React.ReactNode;
type: "input" | "select" | "numerical" | "multi-select";
type: "input" | "select" | "numerical" | "multi-select" | "budget-duration";
options?: Array<{ label: string; value: string }>;
rules?: any[];
step?: number;
@ -155,6 +157,8 @@ const MemberModal = <T extends BaseMember>({
allowClear
/>
);
case "budget-duration":
return <BudgetDurationDropdown />;
default:
return null;
}

13
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View File

@ -388,6 +388,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
max_budget_in_team: values.max_budget_in_team,
tpm_limit: values.tpm_limit,
rpm_limit: values.rpm_limit,
budget_duration: values.budget_duration,
allowed_models: values.allowed_models,
};
MessageManager.destroy(); // Remove all existing toasts
@ -1689,6 +1690,18 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
min: 0,
placeholder: "Budget limit for this member within this team",
},
{
name: "budget_duration",
label: (
<span>
Budget Reset Period{" "}
<Tooltip title="How often this member's budget resets within the team. Leave unset and the budget never resets.">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
),
type: "budget-duration" as const,
},
{
name: "tpm_limit",
label: (

1
ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx
View File

@ -210,6 +210,7 @@ export default function TeamMemberTab({
max_budget_in_team: membership?.litellm_budget_table?.max_budget || null,
tpm_limit: membership?.litellm_budget_table?.tpm_limit || null,
rpm_limit: membership?.litellm_budget_table?.rpm_limit || null,
budget_duration: membership?.litellm_budget_table?.budget_duration || null,
allowed_models: membership?.litellm_budget_table?.allowed_models || [],
};
setSelectedEditMember(enhancedMember);

7
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View File

@ -27574,6 +27574,11 @@ export interface components {
* @description List of models this team member can access. Pass an empty list to remove per-member model restrictions.
*/
allowed_models?: string[] | null;
/**
* Budget Duration
* @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.
*/
budget_duration?: string | null;
/** Max Budget In Team */
max_budget_in_team?: number | null;
/** Role */
@ -27599,6 +27604,8 @@ export interface components {
TeamMemberUpdateResponse: {
/** Allowed Models */
allowed_models?: string[] | null;
/** Budget Duration */
budget_duration?: string | null;
/** Max Budget In Team */
max_budget_in_team?: number | null;
/** Rpm Limit */