From f31d059aa39e7bf982e7580ac6c78f1eb40459a5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 6 Jun 2026 17:24:55 -0700 Subject: [PATCH] feat(ui): add budget duration to edit team member form (#29717) * 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 --- litellm/proxy/_types.py | 5 + .../management_endpoints/common_utils.py | 161 +++---- .../management_endpoints/team_endpoints.py | 55 ++- .../test_upsert_budget_membership.py | 417 ++++++++---------- .../proxy/test_team_member_update.py | 142 +++++- .../src/components/networking.tsx | 18 +- .../src/components/team/EditMembership.tsx | 8 +- .../src/components/team/TeamInfo.tsx | 13 + .../src/components/team/TeamMemberTab.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 + 10 files changed, 512 insertions(+), 315 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e5d7093306..88be567e59 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index dc27e87726..31d831d773 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae7da0d29f..a3ad7a9ea8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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, ) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index f4bf0d7b2b..e9b4f11e89 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -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() diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 6561ec9e7f..352c68d491 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -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() diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 80f4d0b57f..d303734dff 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index af5f8631ae..16e4ec58dd 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -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 = ({ 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 = ({ 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 = ({ allowClear /> ); + case "budget-duration": + return ; default: return null; } diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 4943b1fbac..602490a0c9 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -388,6 +388,7 @@ const TeamInfoView: React.FC = ({ 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 = ({ min: 0, placeholder: "Budget limit for this member within this team", }, + { + name: "budget_duration", + label: ( + + Budget Reset Period{" "} + + + + + ), + type: "budget-duration" as const, + }, { name: "tpm_limit", label: ( diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 5a290a6f6d..e2f108dcbf 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -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); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d14ec8af56..936747acd0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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 */