From 93b3692ad1230840dbb132df38a9958b0fbd79e6 Mon Sep 17 00:00:00 2001 From: RoyVivat Date: Fri, 3 Apr 2026 15:16:57 -0700 Subject: [PATCH] fix: allow clearing budget_duration via UI; sync schema copies Frontend: remove the != null guard in teamMemberUpdateCall so that explicitly selecting 'No reset (unlimited)' sends budget_duration: null to the backend (rather than omitting the field entirely). Backend: add clear_budget_duration flag to _upsert_budget_and_membership. The team_member_update endpoint sets this flag when budget_duration is present in model_fields_set but null, allowing the update path to explicitly null out budget_duration and budget_reset_at on the existing budget row. Also sync litellm/proxy/schema.prisma and litellm-proxy-extras schema to root (missing total_spend column + alignment drift). --- .../litellm_proxy_extras/schema.prisma | 11 +-- .../management_endpoints/common_utils.py | 7 ++ .../management_endpoints/team_endpoints.py | 6 ++ litellm/proxy/schema.prisma | 9 ++- .../management_endpoints/test_common_utils.py | 72 +++++++++++++++++++ .../src/components/networking.tsx | 3 +- 6 files changed, 97 insertions(+), 11 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d8d6015ce87..46be6b31e1f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -321,11 +321,12 @@ model LiteLLM_MCPServerTable { byok_description String[] @default([]) byok_api_key_help_url String? source_url String? - approval_status String? @default("active") - submitted_by String? - submitted_at DateTime? - reviewed_at DateTime? - review_notes String? + // BYOM submission lifecycle + approval_status String? @default("active") + submitted_by String? + submitted_at DateTime? + reviewed_at DateTime? + review_notes String? @@index([approval_status]) } diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index ca3120a5cad..8937abca7af 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -355,6 +355,7 @@ async def _upsert_budget_and_membership( 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 @@ -368,6 +369,9 @@ async def _upsert_budget_and_membership( 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, 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. @@ -404,6 +408,9 @@ async def _upsert_budget_and_membership( # 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, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ed75541d357..0d9bed8b224 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2446,6 +2446,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, @@ -2457,6 +2462,7 @@ async def team_member_update( 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/schema.prisma b/litellm/proxy/schema.prisma index 81125976245..46be6b31e1f 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -595,11 +595,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) - total_spend Float @default(0.0) - budget_id String? + user_id String + team_id String + 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/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/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index be509835d57..73e89961198 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3893,7 +3893,8 @@ 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) { + if (formValues.budget_duration !== undefined) { + // Explicitly include null so the backend can clear a previously-set duration. requestBody.budget_duration = formValues.budget_duration; }