diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql new file mode 100644 index 00000000000..f06c138284a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql @@ -0,0 +1,4 @@ +-- Add total_spend column to LiteLLM_TeamMembership +-- Tracks lifetime (never-zeroed) spend for a user within a team, +-- independent of the current-period spend that resets periodically. +ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fce95465b55..12beefd68c6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -608,9 +608,10 @@ 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) + 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/litellm/proxy/_types.py b/litellm/proxy/_types.py index 441b3b836a1..039705dd026 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1943,6 +1943,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=()) @@ -3622,6 +3623,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]: @@ -3750,6 +3752,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 bcfaed24398..8bcff907ff4 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -52,13 +52,17 @@ 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 # Reset spend counters for affected team members. # Reset Redis directly so a transient failure doesn't leave stale @@ -658,18 +662,25 @@ 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( - seconds=duration_s - ) + # Normal roll-forward: advance from the previous reset time. + # Use a while loop so proxies that were down for multiple + # periods catch up in one pass rather than firing the + # spend-zero operation repeatedly over many scheduler ticks. + while budget.budget_reset_at <= current_time: + budget.budget_reset_at = budget.budget_reset_at + timedelta( + seconds=duration_s + ) except Exception as e: verbose_proxy_logger.exception( "Error resetting budget_reset_at for budget: %s. Item: %s", e, budget diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 241b66bc0ae..76ffcb325df 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1295,7 +1295,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 41a98fa4ad9..3bcfa6e470e 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -173,12 +173,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..b2355697d2d 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -354,6 +354,8 @@ 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, + clear_budget_duration: bool = False, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -366,11 +368,15 @@ 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') + clear_budget_duration: When True, explicitly clears budget_duration (and + budget_reset_at) on an existing budget. Distinct from budget_duration=None + which means "not supplied / don't touch". - 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 and not clear_budget_duration: # 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,11 +395,34 @@ 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 + ) + + if existing_budget_id is not None: + # Update in-place: patch only the fields that were supplied so we don't + # overwrite fields the caller didn't touch (e.g. keep max_budget when + # only budget_duration changes). Exclude created_by — that's set once. + update_data = {k: v for k, v in create_data.items() if k != "created_by"} + if clear_budget_duration: + update_data["budget_duration"] = None + update_data["budget_reset_at"] = None + await tx.litellm_budgettable.update( + where={"budget_id": existing_budget_id}, + data=update_data, + ) + budget_id_to_connect = existing_budget_id + else: + new_budget = await tx.litellm_budgettable.create( + data=create_data, + include={"team_membership": True}, + ) + budget_id_to_connect = new_budget.budget_id - new_budget = await tx.litellm_budgettable.create( - data=create_data, - include={"team_membership": True}, - ) # upsert the team membership with the new/updated budget await tx.litellm_teammembership.upsert( where={ @@ -407,12 +436,12 @@ async def _upsert_budget_and_membership( "user_id": user_id, "team_id": team_id, "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, + "connect": {"budget_id": budget_id_to_connect}, }, }, "update": { "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, + "connect": {"budget_id": budget_id_to_connect}, }, }, }, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e4e0b64af59..205d811ff32 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -231,7 +231,7 @@ class TeamMemberBudgetHandler: updated_kv["metadata"] = {} updated_kv["metadata"]["team_member_budget_id"] = budget_row.budget_id - else: # budget does not exist + else: # budget does not exist — newly creating the template updated_kv = await TeamMemberBudgetHandler.create_team_member_budget_table( data=team_table, new_team_data_json=updated_kv, @@ -242,6 +242,27 @@ class TeamMemberBudgetHandler: team_member_budget_duration=team_member_budget_duration, ) + # Wire up existing members whose budget_id is null to the + # newly-created template budget so they immediately inherit the + # reset schedule. Only do this on creation (not on every update) + # to avoid silently constraining members who were intentionally + # left unlimited. + if team_table.team_id is not None: + new_budget_id: Optional[str] = ( + updated_kv.get("metadata", {}).get("team_member_budget_id") + ) + if new_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": new_budget_id}, + ) + # Remove team member fields from updated_kv TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) return updated_kv @@ -2427,6 +2448,11 @@ async def team_member_update( break ### upsert new budget + # budget_duration=None means "not supplied" by default; we need to distinguish + # that from the user explicitly sending null to clear a previously-set value. + clear_budget_duration = ( + "budget_duration" in data.model_fields_set and data.budget_duration is None + ) async with prisma_client.db.tx() as tx: await _upsert_budget_and_membership( tx=tx, @@ -2437,6 +2463,8 @@ 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, + clear_budget_duration=clear_budget_duration, ) ### update team member role diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9738ae4f1a2..2dac9ead307 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6273,6 +6273,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 fce95465b55..12beefd68c6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -608,9 +608,10 @@ 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) + 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/schema.prisma b/schema.prisma index fce95465b55..12beefd68c6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -608,9 +608,10 @@ 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) + 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..396df6b177c 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,185 @@ 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_reset_at_date_multi_period_catchup(): + """ + When budget_reset_at is non-null but multiple periods have elapsed (e.g. + the proxy was down for 3 weeks and the period is 7 days), the while loop + must advance the timestamp all the way to the next future date in a single + call rather than requiring the scheduler to fire 3 times. + """ + from litellm.proxy._types import LiteLLM_BudgetTableFull + + now = datetime.now(timezone.utc) + # reset_at is 22 days in the past; 7d period → 3 full periods have elapsed + previous_reset = now - timedelta(days=22) + + budget = LiteLLM_BudgetTableFull( + budget_id="b-multi", + budget_duration="7d", + budget_reset_at=previous_reset, + created_at=now - timedelta(days=30), + ) + + result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) + + assert result.budget_reset_at > now, "reset must be in the future after catch-up" + # Should be previous_reset + 28d (4 periods) = now + 6d + expected = previous_reset + timedelta(days=28) + assert result.budget_reset_at == expected, ( + f"Expected {expected}, got {result.budget_reset_at}. " + "Multi-period catch-up must advance through all elapsed periods in one call." + ) + + +@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_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 8b7b5a6fb7a..00557200d0b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -486,3 +486,75 @@ class TestSetObjectMetadataField: ): _set_object_metadata_field(team, "model_rpm_limit", {"x": 1}) assert team.metadata == {"model_rpm_limit": {"x": 1}} + + +class TestUpsertBudgetAndMembershipClearDuration: + """ + Tests for the clear_budget_duration flag in _upsert_budget_and_membership. + Regression for the bug where selecting 'No reset (unlimited)' in the UI + could not clear a previously-set budget_duration because null was filtered + out before reaching the backend. + """ + + def _make_tx(self): + tx = MagicMock() + tx.litellm_budgettable = MagicMock() + tx.litellm_budgettable.update = AsyncMock() + tx.litellm_budgettable.create = AsyncMock( + return_value=MagicMock(budget_id="new-budget-id") + ) + tx.litellm_teammembership = MagicMock() + tx.litellm_teammembership.update = AsyncMock() + tx.litellm_teammembership.upsert = AsyncMock() + return tx + + def _make_auth(self): + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + + @pytest.mark.asyncio + async def test_clear_budget_duration_sets_null_on_existing_budget(self): + """clear_budget_duration=True must null out budget_duration and budget_reset_at.""" + from litellm.proxy.management_endpoints.common_utils import ( + _upsert_budget_and_membership, + ) + + tx = self._make_tx() + await _upsert_budget_and_membership( + tx=tx, + team_id="team-1", + user_id="user-1", + max_budget=None, + existing_budget_id="existing-budget-id", + user_api_key_dict=self._make_auth(), + budget_duration=None, + clear_budget_duration=True, + ) + + tx.litellm_budgettable.update.assert_awaited_once() + update_data = tx.litellm_budgettable.update.call_args.kwargs["data"] + assert update_data["budget_duration"] is None + assert update_data["budget_reset_at"] is None + + @pytest.mark.asyncio + async def test_no_clear_flag_does_not_touch_budget_duration(self): + """Without clear_budget_duration, omitting budget_duration must not overwrite it.""" + from litellm.proxy.management_endpoints.common_utils import ( + _upsert_budget_and_membership, + ) + + tx = self._make_tx() + await _upsert_budget_and_membership( + tx=tx, + team_id="team-1", + user_id="user-1", + max_budget=50.0, + existing_budget_id="existing-budget-id", + user_api_key_dict=self._make_auth(), + budget_duration=None, + clear_budget_duration=False, + ) + + tx.litellm_budgettable.update.assert_awaited_once() + update_data = tx.litellm_budgettable.update.call_args.kwargs["data"] + assert "budget_duration" not in update_data + assert "budget_reset_at" not in update_data 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 232c698603a..5098a9a74ad 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1627,6 +1627,7 @@ async def test_upsert_team_member_budget_table_no_existing_budget(): team_table = MagicMock(spec=LiteLLM_TeamTable) team_table.metadata = {} + team_table.team_id = "test_team_id" team_table.team_alias = "Test Team" team_table.budget_duration = None @@ -1639,12 +1640,22 @@ async def test_upsert_team_member_budget_table_no_existing_budget(): mock_budget_response = MagicMock() mock_budget_response.budget_id = "new_budget_456" - with patch( - "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget", - new_callable=AsyncMock - ) as mock_new_budget: - mock_new_budget.return_value = mock_budget_response + mock_db = MagicMock() + mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 0}) + mock_prisma = MagicMock() + mock_prisma.db = mock_db + with ( + patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget", + new_callable=AsyncMock, + return_value=mock_budget_response, + ) as mock_new_budget, + patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), + ): result = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=team_table, user_api_key_dict=mock_user_api_key_dict, @@ -1661,6 +1672,125 @@ 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_on_create(): + """ + Regression test: when a team budget is being CREATED for the first time, + existing team members whose budget_id is null must be assigned the new + template budget so that their 'Next Budget Reset' column reflects the reset time. + + Previously, members added before the team's budget duration was configured had + budget_id=null. Creating the template budget 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" + ) + + # No existing budget_id — triggers the create path + team_table = MagicMock(spec=LiteLLM_TeamTable) + team_table.metadata = {} + team_table.team_id = "team-xyz" + team_table.team_alias = "Test Team" + team_table.budget_duration = None + + updated_kv: dict = {"team_id": "team-xyz"} + + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "new-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.new_budget", + new_callable=AsyncMock, + return_value=mock_budget_response, + ), + patch( + "litellm.proxy.proxy_server.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"] == "new-template-budget-abc" + + +@pytest.mark.asyncio +async def test_upsert_team_member_budget_table_update_does_not_wire_up_null_members(): + """ + When the team already has a template budget and the settings are UPDATED, + members with budget_id=null must NOT be silently assigned the budget. + Those members were intentionally left unlimited; assigning a budget on + every settings save would be a backwards-incompatible behaviour change. + """ + 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" + ) + + # Existing budget_id — triggers the update path + 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": 0}) + + 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.proxy_server.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 NOT have been called — no silent budget assignment on update + mock_db.litellm_teammembership.update_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_update_team_with_team_member_budget_duration(): """ @@ -6499,6 +6629,84 @@ async def test_create_team_member_budget_table_with_duration(): assert result["metadata"]["team_member_budget_id"] == "budget-abc" +@pytest.mark.asyncio +async def test_upsert_budget_and_membership_updates_existing_budget(): + """ + When existing_budget_id is provided, _upsert_budget_and_membership must UPDATE + the existing budget row rather than creating a new one — otherwise each call + silently orphans the previous LiteLLM_BudgetTable row. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + + mock_tx = MagicMock() + mock_tx.litellm_budgettable.update = AsyncMock() + mock_tx.litellm_budgettable.create = AsyncMock() + mock_tx.litellm_teammembership.upsert = AsyncMock() + + await _upsert_budget_and_membership( + tx=mock_tx, + team_id="team-1", + user_id="user-1", + max_budget=50.0, + existing_budget_id="budget-existing-123", + user_api_key_dict=mock_user, + ) + + # Must update, not create + mock_tx.litellm_budgettable.update.assert_awaited_once() + mock_tx.litellm_budgettable.create.assert_not_awaited() + + # Membership must be connected to the existing budget id + upsert_call = mock_tx.litellm_teammembership.upsert.call_args + connect_id = upsert_call.kwargs["data"]["update"]["litellm_budget_table"]["connect"]["budget_id"] + assert connect_id == "budget-existing-123" + + +@pytest.mark.asyncio +async def test_upsert_budget_and_membership_creates_new_budget_when_none(): + """ + When existing_budget_id is None, _upsert_budget_and_membership must CREATE + a new budget row and connect it to the membership. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + + mock_new_budget = MagicMock() + mock_new_budget.budget_id = "budget-new-456" + + mock_tx = MagicMock() + mock_tx.litellm_budgettable.update = AsyncMock() + mock_tx.litellm_budgettable.create = AsyncMock(return_value=mock_new_budget) + mock_tx.litellm_teammembership.upsert = AsyncMock() + + await _upsert_budget_and_membership( + tx=mock_tx, + team_id="team-1", + user_id="user-1", + max_budget=100.0, + existing_budget_id=None, + user_api_key_dict=mock_user, + ) + + # Must create, not update + mock_tx.litellm_budgettable.create.assert_awaited_once() + mock_tx.litellm_budgettable.update.assert_not_awaited() + + # Membership must be connected to the newly created budget id + upsert_call = mock_tx.litellm_teammembership.upsert.call_args + connect_id = upsert_call.kwargs["data"]["update"]["litellm_budget_table"]["connect"]["budget_id"] + assert connect_id == "budget-new-456" + + # --------------------------------------------------------------------------- # Tests for _batch_resolve_access_group_resources # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 28f8d308de7..afb63f1502f 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3759,6 +3759,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) => { @@ -3898,6 +3899,10 @@ export const teamMemberUpdateCall = async ( if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) { requestBody.rpm_limit = formValues.rpm_limit; } + if (formValues.budget_duration !== undefined) { + // Explicitly include null so the backend can clear a previously-set duration. + 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 6308be65860..09790427241 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -57,6 +57,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; @@ -66,6 +67,7 @@ export interface TeamMembership { rpm_limit: number | null; model_max_budget: Record | null; budget_duration: string | null; + budget_reset_at: string | null; }; } @@ -381,6 +383,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, }; MessageManager.destroy(); // Remove all existing toasts @@ -1483,6 +1486,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);