From bd72f0b42f6803d3494bb51d17ce3374219aa4a9 Mon Sep 17 00:00:00 2001 From: RoyVivat Date: Wed, 1 Apr 2026 12:46:31 -0700 Subject: [PATCH] Add team member budget columns and fix Next Budget Reset display - Add Spend (Current Period), Total Spend, Team Member Budget, Next Budget Reset, and Rate Limits columns to the team members table - Add total_spend field to LiteLLM_TeamMembership (accumulates across resets; never zeroed by the reset job) - Add budget_reset_at to LiteLLM_BudgetTable Python type and schema - Fix Next Budget Reset showing proxy start time: run ResetBudgetJob once immediately at startup via asyncio.create_task so expired budgets are advanced before the first UI request - Fix update_budget to only recalculate budget_reset_at when duration genuinely changes (avoid clobbering a valid future reset on idempotent startup calls) - Fix _reset_budget_reset_at_date to anchor null-reset budgets to created_at instead of current_time - Add Budget Reset Period selector to Edit Member modal; wire budget_duration through TeamMemberUpdateRequest and _upsert_budget_and_membership so individual member budgets get a budget_reset_at on creation - Wire members with budget_id=null to the team template budget on settings update Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/_types.py | 6 + .../proxy/common_utils/reset_budget_job.py | 47 ++-- litellm/proxy/db/db_spend_update_writer.py | 5 +- .../budget_management_endpoints.py | 34 ++- .../management_endpoints/common_utils.py | 13 +- .../management_endpoints/team_endpoints.py | 21 ++ litellm/proxy/proxy_server.py | 6 + litellm/proxy/schema.prisma | 9 +- .../common_utils/test_reset_budget_job.py | 254 +++++++++++++++++- .../test_upsert_budget_membership.py | 102 +++++++ .../test_budget_endpoints.py | 254 ++++++++++++++++++ .../test_team_endpoints.py | 61 +++++ .../src/components/networking.tsx | 4 + .../src/components/team/EditMembership.tsx | 4 +- .../src/components/team/TeamInfo.tsx | 21 ++ .../src/components/team/TeamMemberTab.tsx | 55 +++- 16 files changed, 857 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 91a953c217e..5e0da1900d2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1934,6 +1934,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None model_max_budget: Optional[dict] = None budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None model_config = ConfigDict(protected_namespaces=()) @@ -3607,6 +3608,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): team_id: str budget_id: Optional[str] = None spend: Optional[float] = 0.0 + total_spend: Optional[float] = 0.0 litellm_budget_table: Optional[LiteLLM_BudgetTable] def safe_get_team_member_rpm_limit(self) -> Optional[int]: @@ -3735,6 +3737,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="Budget reset period for this team member (e.g. '24h', '7d', '30d'). Determines how often their individual spend resets.", + ) class TeamMemberUpdateResponse(MemberUpdateResponse): diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 674214b19e5..88e3cdd2134 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -52,21 +52,21 @@ class ResetBudgetJob: self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): """ - Resets the budget for all LiteLLM Team Members if their budget has expired + Resets the budget for all LiteLLM Team Members if their budget has expired. + Only resets the current-period spend; total_spend is never zeroed so it + accumulates as a lifetime usage counter. """ + budget_ids = [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + if not budget_ids: + return + return await self.prisma_client.db.litellm_teammembership.update_many( - where={ - "budget_id": { - "in": [ - budget.budget_id - for budget in budgets_to_reset - if budget.budget_id is not None - ] - } - }, - data={ - "spend": 0, - }, + where={"budget_id": {"in": budget_ids}}, + data={"spend": 0}, ) async def reset_budget_for_keys_linked_to_budgets( @@ -590,16 +590,19 @@ class ResetBudgetJob: duration_s = duration_in_seconds(duration=budget.budget_duration) - # Fallback for existing budgets that do not have a budget_reset_at date set, ensuring the duration is taken into account - if ( - budget.budget_reset_at is None - and budget.created_at + timedelta(seconds=duration_s) > current_time - ): - budget.budget_reset_at = budget.created_at + timedelta( - seconds=duration_s - ) + if budget.budget_reset_at is None: + # Budget was created without a reset time. Anchor the first + # period to created_at so the reset aligns with when the + # budget was originally set up, not the current wall-clock + # time (which would be proxy-start time when this job runs + # at startup and produce a confusing "resets at startup" display). + anchor = budget.created_at + timedelta(seconds=duration_s) + while anchor <= current_time: + anchor += timedelta(seconds=duration_s) + budget.budget_reset_at = anchor else: - budget.budget_reset_at = current_time + timedelta( + # Normal roll-forward: advance from the previous reset time. + budget.budget_reset_at = budget.budget_reset_at + timedelta( seconds=duration_s ) except Exception as e: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a305d5be1e6..0db119b5337 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1251,7 +1251,10 @@ class DBSpendUpdateWriter: batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists where={"team_id": team_id, "user_id": user_id}, - data={"spend": {"increment": response_cost}}, + data={ + "spend": {"increment": response_cost}, + "total_spend": {"increment": response_cost}, + }, ) # Transaction succeeded, break out of retry loop break diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 20c7f9ec412..eac239c7981 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -161,12 +161,38 @@ async def update_budget( except ValueError as e: raise HTTPException(status_code=400, detail={"error": str(e)}) + update_data: dict = { + **budget_obj.model_dump(exclude_unset=True), + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + + # Recalculate budget_reset_at only when budget_duration is genuinely changing. + # We compare against the current DB value so that idempotent calls at proxy + # startup (e.g. _update_default_team_member_budget) do not overwrite a + # correctly-set future budget_reset_at with "now + duration". + # We check budget_obj.budget_duration directly (not model_fields_set) because + # internal callers set fields via attribute assignment after construction, which + # does not reliably update model_fields_set. + if budget_obj.budget_duration is not None and "budget_reset_at" not in update_data: + existing = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": budget_obj.budget_id} + ) + existing_duration = existing.budget_duration if existing is not None else None + duration_changed = existing_duration != budget_obj.budget_duration + reset_at_missing = ( + existing is None or existing.budget_reset_at is None + ) + if duration_changed or reset_at_missing: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + update_data["budget_duration"] = budget_obj.budget_duration + update_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=budget_obj.budget_duration + ) + response = await prisma_client.db.litellm_budgettable.update( where={"budget_id": budget_obj.budget_id}, - data={ - **budget_obj.model_dump(exclude_unset=True), # type: ignore - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - }, # type: ignore + data=update_data, # type: ignore ) return response diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index efc42d3355c..98b59adbeb5 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -354,6 +354,7 @@ async def _upsert_budget_and_membership( user_api_key_dict: UserAPIKeyAuth, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -366,11 +367,12 @@ async def _upsert_budget_and_membership( 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 + budget_duration: Budget reset period (e.g. '24h', '7d', '30d') - If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership. + If max_budget, tpm_limit, rpm_limit, and budget_duration 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. """ - if max_budget is None and tpm_limit is None and rpm_limit is None: + if max_budget is None and tpm_limit is None and rpm_limit is None and budget_duration 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}}, @@ -389,6 +391,13 @@ async def _upsert_budget_and_membership( create_data["tpm_limit"] = tpm_limit if rpm_limit is not None: create_data["rpm_limit"] = rpm_limit + if budget_duration is not None: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + create_data["budget_duration"] = budget_duration + create_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=budget_duration + ) 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 3643373be65..b879fd031ac 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -240,6 +240,26 @@ class TeamMemberBudgetHandler: team_member_budget_duration=team_member_budget_duration, ) + # Wire up any existing members whose budget_id is null so they + # immediately see the template budget (and its budget_reset_at). + # This handles members who were added before the team's budget + # duration was ever configured. + if team_table.team_id is not None: + final_budget_id: Optional[str] = ( + updated_kv.get("metadata", {}).get("team_member_budget_id") + ) + if final_budget_id is not None: + from litellm.proxy.proxy_server import prisma_client as _prisma_client + + if _prisma_client is not None: + await _prisma_client.db.litellm_teammembership.update_many( + where={ + "team_id": team_table.team_id, + "budget_id": None, + }, + data={"budget_id": final_budget_id}, + ) + # Remove team member fields from updated_kv TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) return updated_kv @@ -2435,6 +2455,7 @@ async def team_member_update( user_api_key_dict=user_api_key_dict, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + budget_duration=data.budget_duration, ) ### update team member role diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9c29927c5cb..cad1f0a8010 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5930,6 +5930,12 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + # Run once immediately at startup so expired budget_reset_at values are + # advanced before the first UI request. Without this, any budget that + # expired exactly at restart time would show a stale past timestamp + # until the first scheduled interval fires (~10 minutes later). + asyncio.create_task(budget_reset_job.reset_budget()) + ### UPDATE SPEND ### scheduler.add_job( update_spend, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 46be6b31e1f..81125976245 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -595,10 +595,11 @@ model LiteLLM_UserNotifications { model LiteLLM_TeamMembership { // Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team - user_id String - team_id String - spend Float @default(0.0) - budget_id String? + user_id String + team_id String + spend Float @default(0.0) + total_spend Float @default(0.0) + budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) @@id([user_id, team_id]) } diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index f975460836a..2625aae0e0d 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -18,10 +18,13 @@ from litellm.proxy.utils import ProxyLogging # Mock classes for testing class MockLiteLLMTeamMembership: + def __init__(self): + self.update_many_calls: List[Dict[str, Any]] = [] + async def update_many( self, where: Dict[str, Any], data: Dict[str, Any] ) -> Dict[str, Any]: - # Mock the update_many method for litellm_teammembership + self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -41,6 +44,10 @@ class MockDB: self.litellm_teammembership = MockLiteLLMTeamMembership() self.litellm_verificationtoken = MockLiteLLMVerificationToken() + def reset(self): + self.litellm_teammembership = MockLiteLLMTeamMembership() + self.litellm_verificationtoken = MockLiteLLMVerificationToken() + class MockPrismaClient: def __init__(self): @@ -444,6 +451,101 @@ def test_reset_budget_for_keys_linked_to_budgets_empty( assert len(calls) == 0 +def test_reset_budget_for_team_members_resets_spend( + reset_budget_job, mock_prisma_client +): + """ + Test that reset_budget_for_litellm_team_members resets spend to 0 + for all team memberships linked to the expired budgets. + """ + now = datetime.now(timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 50.0, + "budget_duration": "1mo", + "budget_reset_at": now, + "budget_id": "team-member-budget-1", + "created_at": now - timedelta(days=30), + }, + ) + + budgets_to_reset = [test_budget] + + asyncio.run( + reset_budget_job.reset_budget_for_litellm_team_members( + budgets_to_reset=budgets_to_reset + ) + ) + + calls = mock_prisma_client.db.litellm_teammembership.update_many_calls + assert len(calls) == 1 + + call = calls[0] + # spend must be zeroed out for the matching budget IDs + assert call["where"]["budget_id"] == {"in": ["team-member-budget-1"]} + assert call["data"]["spend"] == 0 + + +def test_reset_budget_for_team_members_preserves_total_spend( + reset_budget_job, mock_prisma_client +): + """ + Test that reset_budget_for_litellm_team_members only resets the current-period + spend and does NOT reset total_spend. total_spend must accumulate across all + budget periods to give a lifetime view of member usage. + """ + now = datetime.now(timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 50.0, + "budget_duration": "1mo", + "budget_reset_at": now, + "budget_id": "team-member-budget-2", + "created_at": now - timedelta(days=30), + }, + ) + + budgets_to_reset = [test_budget] + + asyncio.run( + reset_budget_job.reset_budget_for_litellm_team_members( + budgets_to_reset=budgets_to_reset + ) + ) + + calls = mock_prisma_client.db.litellm_teammembership.update_many_calls + assert len(calls) == 1 + + call = calls[0] + # total_spend must NOT be present in the reset data — it must never be zeroed + assert "total_spend" not in call["data"], ( + "total_spend should not be reset during a budget period reset; " + "it must accumulate across all periods" + ) + assert call["data"]["spend"] == 0 + + +def test_reset_budget_for_team_members_empty_budgets( + reset_budget_job, mock_prisma_client +): + """ + Test that reset_budget_for_litellm_team_members does nothing when given + an empty list of budgets (e.g. no budget periods have expired). + """ + asyncio.run( + reset_budget_job.reset_budget_for_litellm_team_members(budgets_to_reset=[]) + ) + + calls = mock_prisma_client.db.litellm_teammembership.update_many_calls + assert len(calls) == 0 + + def test_budget_table_reset_also_resets_linked_keys( reset_budget_job, mock_prisma_client ): @@ -479,3 +581,153 @@ def test_budget_table_reset_also_resets_linked_keys( ) assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]} assert calls[0]["data"]["spend"] == 0 + + +# --------------------------------------------------------------------------- +# Tests for _reset_budget_reset_at_date +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_reset_budget_reset_at_date_normal_rollforward(): + """ + When budget_reset_at is set and has expired, roll it forward by one + duration interval from the PREVIOUS reset time (not from now). + This ensures periods stay aligned regardless of when the job runs. + """ + from litellm.proxy._types import LiteLLM_BudgetTableFull + + now = datetime.now(timezone.utc) + previous_reset = now - timedelta(days=1) # expired yesterday + + budget = LiteLLM_BudgetTableFull( + budget_id="b1", + budget_duration="7d", + budget_reset_at=previous_reset, + created_at=now - timedelta(days=30), + ) + + result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) + + expected = previous_reset + timedelta(days=7) + assert result.budget_reset_at == expected, ( + "budget_reset_at should advance by exactly one duration from the previous reset" + ) + assert result.budget_reset_at > now, "new reset must be in the future" + + +@pytest.mark.asyncio +async def test_reset_budget_reset_at_date_none_reset_at_not_yet_expired(): + """ + When budget_reset_at is None but created_at + duration is still in the + future, anchor to created_at + duration (the original planned reset time). + """ + from litellm.proxy._types import LiteLLM_BudgetTableFull + + now = datetime.now(timezone.utc) + created_at = now - timedelta(days=3) # created 3 days ago + # duration is 7d → first reset would be created_at + 7d = 4 days from now + + budget = LiteLLM_BudgetTableFull( + budget_id="b2", + budget_duration="7d", + budget_reset_at=None, + created_at=created_at, + ) + + result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) + + expected = created_at + timedelta(days=7) + assert result.budget_reset_at == expected + assert result.budget_reset_at > now + + +@pytest.mark.asyncio +async def test_reset_budget_reset_at_date_none_reset_at_already_past(): + """ + Regression: when budget_reset_at is None AND created_at + duration is + already in the past (old budget that was never given a reset time), the + new reset_at must be computed by advancing from created_at in duration + steps until it's in the future — NOT set to current_time + duration. + + Setting it to current_time + duration was the bug: it made the reset + display show "proxy start time + duration" because this path runs at + startup via the ResetBudgetJob scheduler. + """ + from litellm.proxy._types import LiteLLM_BudgetTableFull + + now = datetime.now(timezone.utc) + # Budget was created 30 days ago with a 7-day duration → 4 full periods + # have elapsed; created_at + 7d, +14d, +21d, +28d are all in the past. + created_at = now - timedelta(days=30) + + budget = LiteLLM_BudgetTableFull( + budget_id="b3", + budget_duration="7d", + budget_reset_at=None, + created_at=created_at, + ) + + result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) + + assert result.budget_reset_at is not None + assert result.budget_reset_at > now, "reset must be in the future" + + # Must be anchored to created_at, not to `now`. + # The correct next reset is created_at + N*7d where N is the smallest + # integer that puts it past `now`. For 30 days elapsed: N=5 → +35d. + expected = created_at + timedelta(days=35) + assert result.budget_reset_at == expected, ( + f"Expected {expected}, got {result.budget_reset_at}. " + "reset_at must be anchored to created_at, not proxy-start time." + ) + + # Verify it is NOT simply now + duration (the old buggy behaviour). + buggy_value = now + timedelta(days=7) + assert result.budget_reset_at != pytest.approx(buggy_value, abs=timedelta(minutes=1)), ( + "reset_at must not be set to current_time + duration (proxy-start-time bug)" + ) + + +@pytest.mark.asyncio +async def test_reset_budget_at_startup_advances_expired_budget_reset_at( + reset_budget_job, mock_prisma_client +): + """ + Regression: on proxy restart, any budget whose budget_reset_at has just + expired (e.g. budget created 24 h ago, reset_at = proxy_start_time) must be + advanced by the immediate reset_budget() call that now runs at startup. + + Without the startup call, a budget that expired at the exact moment the proxy + restarted would show a past timestamp in the UI for up to ~10 minutes (the + first scheduled interval), even if the user checked 50 minutes later. + """ + now = datetime.now(timezone.utc) + proxy_start_time = now # simulated proxy start time + + # Budget was created 24 h ago; reset_at is exactly now (proxy_start_time) + expired_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "24h", + "budget_reset_at": proxy_start_time, # expired right now + "budget_id": "template-budget-1", + "created_at": proxy_start_time - timedelta(hours=24), + }, + ) + + mock_prisma_client.data["budget"] = [expired_budget] + + # Simulate what the startup asyncio.create_task fires + await reset_budget_job.reset_budget_for_litellm_budget_table() + + # The budget should have been advanced in DB + updated_budgets = mock_prisma_client.updated_data.get("budget", []) + assert len(updated_budgets) == 1, "Expected the expired budget to be updated" + + updated = updated_budgets[0] + assert updated.budget_reset_at > proxy_start_time, ( + "After startup reset, budget_reset_at must be in the future, " + "not equal to proxy start time" + ) 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 a36dc7ff2e3..6d0736c25e5 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 @@ -316,3 +316,105 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): }, }, ) + + +# TEST: budget_duration sets budget_reset_at on the new budget +@pytest.mark.asyncio +async def test_upsert_with_budget_duration_sets_reset_at(mock_tx, fake_user): + """ + When budget_duration is provided, the new individual member budget should + have both budget_duration and budget_reset_at written so that the + Next Budget Reset column shows a real future time instead of '-'. + """ + from datetime import datetime, timezone + from unittest.mock import patch + + fake_reset_time = datetime(2026, 5, 1, 0, 0, 0, tzinfo=timezone.utc) + + with patch( + "litellm.proxy.common_utils.timezone_utils.get_budget_reset_time", + return_value=fake_reset_time, + ) as mock_get_reset: + await _upsert_budget_and_membership( + mock_tx, + team_id="team-dur", + user_id="user-dur", + max_budget=10.0, + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_duration="24h", + ) + + mock_get_reset.assert_called_once_with(budget_duration="24h") + + mock_tx.litellm_budgettable.create.assert_awaited_once_with( + data={ + "max_budget": 10.0, + "budget_duration": "24h", + "budget_reset_at": fake_reset_time, + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + }, + include={"team_membership": True}, + ) + + +# TEST: budget_duration=None alone does not create a budget (disconnect path) +@pytest.mark.asyncio +async def test_upsert_disconnect_when_only_budget_duration_none(mock_tx, fake_user): + """ + Passing only None values (including budget_duration=None) should still + hit the disconnect path, not attempt to create a budget. + """ + await _upsert_budget_and_membership( + mock_tx, + team_id="team-none", + user_id="user-none", + max_budget=None, + existing_budget_id=None, + user_api_key_dict=fake_user, + tpm_limit=None, + rpm_limit=None, + budget_duration=None, + ) + + mock_tx.litellm_teammembership.update.assert_awaited_once() + mock_tx.litellm_budgettable.create.assert_not_called() + + +# TEST: budget_duration alone (no max_budget/limits) creates a budget +@pytest.mark.asyncio +async def test_upsert_budget_duration_only_creates_budget(mock_tx, fake_user): + """ + Setting only budget_duration (no max_budget, tpm, rpm) should still create + a budget with the duration so that the member gets a reset schedule. + """ + from datetime import datetime, timezone + from unittest.mock import patch + + fake_reset_time = datetime(2026, 5, 8, 0, 0, 0, tzinfo=timezone.utc) + + with patch( + "litellm.proxy.common_utils.timezone_utils.get_budget_reset_time", + return_value=fake_reset_time, + ): + await _upsert_budget_and_membership( + mock_tx, + team_id="team-dur-only", + user_id="user-dur-only", + max_budget=None, + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_duration="7d", + ) + + mock_tx.litellm_budgettable.create.assert_awaited_once_with( + data={ + "budget_duration": "7d", + "budget_reset_at": fake_reset_time, + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + }, + include={"team_membership": True}, + ) + mock_tx.litellm_teammembership.update.assert_not_called() diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index b15b9d622e4..e1b8c1973ba 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -24,6 +24,12 @@ def client_and_mocks(monkeypatch): mock_table = MagicMock() mock_table.create = AsyncMock(side_effect=lambda *, data: data) mock_table.update = AsyncMock(side_effect=lambda *, where, data: {**where, **data}) + # Default find_unique: return a budget whose duration differs from "7d" so + # duration-change tests see a change. Individual tests override this as needed. + mock_table.find_unique = AsyncMock(return_value=types.SimpleNamespace( + budget_duration=None, + budget_reset_at=None, + )) mock_prisma.db = types.SimpleNamespace( litellm_budgettable=mock_table, @@ -265,3 +271,251 @@ async def test_new_budget_invalid_model_max_budget(client_and_mocks, monkeypatch assert resp.status_code in (400, 422), resp.text detail = resp.json()["detail"] assert "model_max_budget" in str(detail) or "dictionary" in str(detail).lower() + + +@pytest.mark.asyncio +async def test_update_budget_duration_recalculates_budget_reset_at(client_and_mocks): + """ + Regression test: changing budget_duration via /budget/update must also update + budget_reset_at to reflect the new period. + + Previously, update_budget() stored the new duration but left the old + budget_reset_at untouched, so the next reset shown in the UI was still in the + past (the old period's reset time). + """ + from datetime import datetime, timezone + + client, _, mock_table = client_and_mocks + + captured_data: dict = {} + + async def capture_update(*, where, data): + captured_data.update(data) + return {**where, **data} + + mock_table.update = AsyncMock(side_effect=capture_update) + + payload = { + "budget_id": "budget_duration_change", + "budget_duration": "7d", + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 200, resp.text + + # budget_reset_at must be present and in the future + assert "budget_reset_at" in captured_data, ( + "budget_reset_at must be recalculated when budget_duration changes" + ) + reset_at = captured_data["budget_reset_at"] + if isinstance(reset_at, str): + reset_at = datetime.fromisoformat(reset_at) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) + assert reset_at > datetime.now(timezone.utc), ( + "budget_reset_at must be in the future after a duration change" + ) + + +@pytest.mark.asyncio +async def test_update_budget_duration_respects_explicit_budget_reset_at(client_and_mocks): + """ + If the caller explicitly provides budget_reset_at alongside budget_duration, + the supplied value must be used as-is and not overwritten. + """ + from datetime import datetime, timezone, timedelta + + client, _, mock_table = client_and_mocks + + captured_data: dict = {} + + async def capture_update(*, where, data): + captured_data.update(data) + return {**where, **data} + + mock_table.update = AsyncMock(side_effect=capture_update) + + explicit_reset = (datetime.now(timezone.utc) + timedelta(days=30)).isoformat() + payload = { + "budget_id": "budget_explicit_reset", + "budget_duration": "7d", + "budget_reset_at": explicit_reset, + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 200, resp.text + + stored = captured_data.get("budget_reset_at") + if isinstance(stored, str): + stored = datetime.fromisoformat(stored) + if isinstance(stored, datetime) and stored.tzinfo is None: + stored = stored.replace(tzinfo=timezone.utc) + + explicit_dt = datetime.fromisoformat(explicit_reset) + # Should be within a second of the explicit value + assert abs((stored - explicit_dt).total_seconds()) < 1, ( + "Explicit budget_reset_at must not be overwritten by the duration recalculation" + ) + + +@pytest.mark.asyncio +async def test_update_budget_duration_via_internal_call_path(client_and_mocks, monkeypatch): + """ + Regression test for the internal call path used by upsert_team_member_budget_table. + + When the team member budget duration is changed via the team settings UI, the + backend constructs a BudgetNewRequest by first calling BudgetNewRequest(budget_id=...) + and then setting budget_duration via attribute assignment AFTER construction: + + budget_request = BudgetNewRequest(budget_id=some_id) + budget_request.budget_duration = "7d" # ← attribute assignment, not __init__ + await update_budget(budget_obj=budget_request, ...) + + In Pydantic v2 without validate_assignment=True, attribute assignment does NOT + reliably update model_fields_set. The previous fix checked model_fields_set and + therefore silently skipped the budget_reset_at recalculation on this path. + + The new fix checks budget_obj.budget_duration directly and writes budget_reset_at + explicitly into the update dict, bypassing model_fields_set entirely. + """ + import litellm.proxy.proxy_server as ps + from datetime import datetime, timezone + from litellm.proxy._types import BudgetNewRequest, UserAPIKeyAuth, LitellmUserRoles + from litellm.proxy.management_endpoints.budget_management_endpoints import update_budget + + import types as _types + + _, mock_prisma, mock_table = client_and_mocks + + # find_unique: existing budget has no duration (None → "7d" is a real change) + mock_table.find_unique = AsyncMock(return_value=_types.SimpleNamespace( + budget_duration=None, + budget_reset_at=None, + )) + + captured_data: dict = {} + + async def capture_update(*, where, data): + captured_data.update(data) + return {**where, **data} + + mock_table.update = AsyncMock(side_effect=capture_update) + + fake_user = UserAPIKeyAuth( + user_id="internal_caller", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + # Simulate exactly what upsert_team_member_budget_table does: + # construct with budget_id only, then set budget_duration via attribute assignment. + budget_request = BudgetNewRequest(budget_id="internal-budget-id") + budget_request.budget_duration = "7d" # attribute assignment — not via __init__ + + await update_budget(budget_obj=budget_request, user_api_key_dict=fake_user) + + assert "budget_reset_at" in captured_data, ( + "budget_reset_at must be recalculated even when budget_duration was set " + "via attribute assignment after BudgetNewRequest construction" + ) + assert "budget_duration" in captured_data, ( + "budget_duration must be present in the DB update" + ) + reset_at = captured_data["budget_reset_at"] + if isinstance(reset_at, str): + reset_at = datetime.fromisoformat(reset_at) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) + assert reset_at > datetime.now(timezone.utc), ( + "budget_reset_at must be a future datetime" + ) + + +@pytest.mark.asyncio +async def test_update_budget_same_duration_does_not_overwrite_reset_at(client_and_mocks): + """ + Regression test: proxy startup calls update_team → upsert_team_member_budget_table + → update_budget with the SAME budget_duration that's already stored in the DB. + + Before this fix, update_budget always recalculated budget_reset_at when + budget_duration was set, so every proxy restart would stamp the template + budget's reset time to "now + duration", making the UI display proxy-start + time as the "Next Budget Reset" for all members. + + When the duration hasn't changed and budget_reset_at is already set to a + future value, update_budget must NOT touch budget_reset_at. + """ + import types as _types + from datetime import datetime, timezone, timedelta + + client, _, mock_table = client_and_mocks + + future_reset = datetime.now(timezone.utc) + timedelta(days=5) + + # Simulate existing budget with the SAME duration and a valid future reset time + mock_table.find_unique = AsyncMock(return_value=_types.SimpleNamespace( + budget_duration="7d", + budget_reset_at=future_reset, + )) + + captured_data: dict = {} + + async def capture_update(*, where, data): + captured_data.update(data) + return {**where, **data} + + mock_table.update = AsyncMock(side_effect=capture_update) + + payload = { + "budget_id": "template-budget-startup", + "budget_duration": "7d", # same as what's already in DB + "max_budget": 50.0, + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 200, resp.text + + assert "budget_reset_at" not in captured_data, ( + "budget_reset_at must NOT be overwritten when budget_duration is unchanged " + "and a valid future reset time already exists — proxy-start-time bug" + ) + + +@pytest.mark.asyncio +async def test_update_budget_changed_duration_recalculates_reset_at(client_and_mocks): + """ + Complementary to the no-op test: when budget_duration genuinely changes + (e.g. user switches from 7d to 1mo in the UI), budget_reset_at must be + recalculated to reflect the new period. + """ + import types as _types + from datetime import datetime, timezone, timedelta + + client, _, mock_table = client_and_mocks + + # Existing budget has "7d"; request changes it to "30d" + mock_table.find_unique = AsyncMock(return_value=_types.SimpleNamespace( + budget_duration="7d", + budget_reset_at=datetime.now(timezone.utc) + timedelta(days=3), + )) + + captured_data: dict = {} + + async def capture_update(*, where, data): + captured_data.update(data) + return {**where, **data} + + mock_table.update = AsyncMock(side_effect=capture_update) + + payload = { + "budget_id": "template-budget-change", + "budget_duration": "30d", # different from DB value + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 200, resp.text + + assert "budget_reset_at" in captured_data, ( + "budget_reset_at must be recalculated when budget_duration changes" + ) + reset_at = captured_data["budget_reset_at"] + if isinstance(reset_at, str): + reset_at = datetime.fromisoformat(reset_at) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) + assert reset_at > datetime.now(timezone.utc) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 366f659bdab..b96f972eb2c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1661,6 +1661,67 @@ async def test_upsert_team_member_budget_table_no_existing_budget(): assert "team_member_budget_duration" not in result +@pytest.mark.asyncio +async def test_upsert_team_member_budget_table_wires_up_null_budget_members(): + """ + Regression test: when team_member_budget_duration is set (create or update path), + existing team members whose budget_id is null must be assigned the template + budget so that their 'Next Budget Reset' column reflects the new reset time. + + Previously, members added before the team's budget duration was configured had + budget_id=null. Changing the duration created/updated the template but left + those memberships disconnected — they never showed a reset date. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" + ) + + team_table = MagicMock(spec=LiteLLM_TeamTable) + team_table.metadata = {"team_member_budget_id": "template-budget-abc"} + team_table.team_id = "team-xyz" + + updated_kv: dict = {"team_id": "team-xyz"} + + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "template-budget-abc" + + mock_db = MagicMock() + mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 2}) + + mock_prisma = MagicMock() + mock_prisma.db = mock_db + + with ( + patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + return_value=mock_budget_response, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints.prisma_client", + mock_prisma, + ), + ): + await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=team_table, + user_api_key_dict=mock_user, + updated_kv=updated_kv, + team_member_budget_duration="7d", + ) + + # update_many must have been called to wire up members with no budget + mock_db.litellm_teammembership.update_many.assert_awaited_once() + call_kwargs = mock_db.litellm_teammembership.update_many.call_args.kwargs + assert call_kwargs["where"]["team_id"] == "team-xyz" + assert call_kwargs["where"]["budget_id"] is None + assert call_kwargs["data"]["budget_id"] == "template-budget-abc" + + @pytest.mark.asyncio async def test_update_team_with_team_member_budget_duration(): """ diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index bb9a0c1e015..61e540cf7eb 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3695,6 +3695,7 @@ export interface Member { max_budget_in_team?: number | null; tpm_limit?: number | null; rpm_limit?: number | null; + budget_duration?: string | null; } export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => { @@ -3834,6 +3835,9 @@ export const teamMemberUpdateCall = async ( if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) { requestBody.rpm_limit = formValues.rpm_limit; } + if (formValues.budget_duration !== undefined && formValues.budget_duration !== null) { + requestBody.budget_duration = formValues.budget_duration; + } console.log("Final request body:", requestBody); diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index e6d9ccad60b..2bea1fbdec0 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -85,8 +85,8 @@ const MemberModal = ({ const formData = Object.entries(values).reduce((acc, [key, value]) => { if (typeof value === "string") { const trimmedValue = value.trim(); - // For empty strings on optional numeric fields, set to null - if (trimmedValue === "" && (key === "max_budget_in_team" || key === "tpm_limit" || key === "rpm_limit")) { + // For empty strings on optional fields, set to null + if (trimmedValue === "" && (key === "max_budget_in_team" || key === "tpm_limit" || key === "rpm_limit" || key === "budget_duration")) { return { ...acc, [key]: null }; } return { ...acc, [key]: trimmedValue }; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index d2ce79580da..2b7fc18ccbf 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -56,6 +56,7 @@ export interface TeamMembership { team_id: string; budget_id: string; spend: number; + total_spend: number | null; litellm_budget_table: { budget_id: string; soft_budget: number | null; @@ -65,6 +66,7 @@ export interface TeamMembership { rpm_limit: number | null; model_max_budget: Record | null; budget_duration: string | null; + budget_reset_at: string | null; }; } @@ -365,6 +367,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 || null, }; message.destroy(); // Remove all existing toasts @@ -1306,6 +1309,24 @@ const TeamInfoView: React.FC = ({ min: 0, placeholder: "Requests per minute limit for this member in this team", }, + { + name: "budget_duration", + label: ( + + Budget Reset Period{" "} + + + + + ), + type: "select" as const, + options: [ + { label: "No reset (unlimited)", value: "" }, + { label: "Daily", value: "24h" }, + { label: "Weekly", value: "7d" }, + { label: "Monthly", value: "30d" }, + ], + }, ], }} /> diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 652d1dbcd91..694c5136e9c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -45,13 +45,27 @@ export default function TeamMemberTab({ return "0"; }; - // Helper function to get spend for a user + // Helper function to get spend for a user (current budget period only) const getUserSpend = (userId: string | null): number | null => { if (!userId) return 0; const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); return membership?.spend || 0; }; + // Helper function to get total (all-time) spend for a user + const getUserTotalSpend = (userId: string | null): number | null => { + if (!userId) return 0; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + return membership?.total_spend ?? 0; + }; + + // Helper function to get the next budget reset datetime for a user + const getUserBudgetResetAt = (userId: string | null): string | null => { + if (!userId) return null; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + return membership?.litellm_budget_table?.budget_reset_at ?? null; + }; + const getUserBudget = (userId: string | null): string | null => { if (!userId) return null; const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); @@ -86,8 +100,8 @@ export default function TeamMemberTab({ { title: ( - Team Member Spend (USD) - + Spend (Current Period) + @@ -97,6 +111,20 @@ export default function TeamMemberTab({ ${formatNumberWithCommas(getUserSpend(record.user_id), 4)} ), }, + { + title: ( + + Total Spend (USD) + + + + + ), + key: "total_spend", + render: (_: unknown, record: Member) => ( + ${formatNumberWithCommas(getUserTotalSpend(record.user_id), 4)} + ), + }, { title: "Team Member Budget (USD)", key: "budget", @@ -109,6 +137,26 @@ export default function TeamMemberTab({ ); }, }, + { + title: ( + + Next Budget Reset + + + + + ), + key: "budget_reset_at", + render: (_: unknown, record: Member) => { + const resetAt = getUserBudgetResetAt(record.user_id); + if (!resetAt) return -; + return ( + + {new Date(resetAt).toLocaleString()} + + ); + }, + }, { title: ( @@ -138,6 +186,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, }; setSelectedEditMember(enhancedMember); setIsEditMemberModalVisible(true);