Merge pull request #40895 from BerriAI/litellm_budget_clear_persistence

fix(ui): persist cleared budgets and reset intervals

enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py

@@ -780,7 +780,10 @@ async def update_project(
 
         # Handle budget updates
         budget_fields = LiteLLM_BudgetTable.model_fields.keys()
-        budget_updates = {k: v for k, v in update_data.items() if k in budget_fields}
+        budget_updates = {
+            **{k: v for k, v in update_data.items() if k in budget_fields},
+            **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}),
+        }
 
         if budget_updates and existing_project.budget_id:
             # Update existing budget

litellm/proxy/management_endpoints/budget_management_endpoints.py

@@ -14,6 +14,7 @@ All /budget management endpoints
 #### BUDGET TABLE MANAGEMENT ####
 import math
 from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Final
 
 from fastapi import APIRouter, Depends, HTTPException
@@ -176,6 +177,10 @@ async def update_budget(
     recomputed_reset_at: Final = (
         {"budget_reset_at": get_budget_reset_time(budget_duration=budget_obj.budget_duration)}
         if budget_obj.budget_duration is not None and "budget_reset_at" not in budget_obj.model_fields_set
+        else MappingProxyType({"budget_reset_at": None})
+        if "budget_duration" in budget_obj.model_fields_set
+        and budget_obj.budget_duration is None
+        and "budget_reset_at" not in budget_obj.model_fields_set
         else {}
     )
 

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -1254,8 +1254,8 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
     fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set()
 
     for k, v in data_json.items():
-        if k == "max_budget":
-            if "max_budget" in fields_set:
+        if k in ("max_budget", "budget_duration"):
+            if k in fields_set:
                 non_default_values[k] = v
         elif k == "model_max_budget":
             if k in fields_set:
@@ -1283,8 +1283,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
         from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
 
         validate_budget_duration(non_default_values["budget_duration"])
-        non_default_values["budget_reset_at"] = get_budget_reset_time(
-            budget_duration=non_default_values["budget_duration"]
+        non_default_values["budget_reset_at"] = (
+            get_budget_reset_time(budget_duration=non_default_values["budget_duration"])
+            if non_default_values["budget_duration"] is not None
+            else None
         )
 
     if "max_budget" not in non_default_values:

litellm/proxy/management_endpoints/tag_management_endpoints.py

@@ -438,6 +438,7 @@ async def update_tag(
             user_api_key_dict=user_api_key_dict,
             prisma_client=prisma_client,
             litellm_proxy_admin_name=litellm_proxy_admin_name,
+            budget_duration_cleared="budget_duration" in tag.model_fields_set and tag.budget_duration is None,
         )
 
         # Get model names for model_info

litellm/proxy/management_helpers/utils.py

@@ -3,6 +3,7 @@
 from collections.abc import Callable, Mapping, MutableMapping, Sequence
 from datetime import datetime
 from functools import wraps
+from types import MappingProxyType
 from typing import Any, Final, Protocol
 
 from fastapi import HTTPException, Request
@@ -180,6 +181,7 @@ async def handle_budget_for_entity(
     user_api_key_dict: UserAPIKeyAuth,
     prisma_client: PrismaClient,
     litellm_proxy_admin_name: str,
+    budget_duration_cleared: bool = False,
 ) -> str | None:
     """
     Common helper to handle budget creation/updates for entities (organizations, tags, etc).
@@ -208,7 +210,14 @@ async def handle_budget_for_entity(
 
     # Extract budget fields from data
     _json_data: Final = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data
-    _budget_data: Final = {k: v for k, v in _json_data.items() if k in budget_params}
+    _budget_data: Final = MappingProxyType(
+        {
+            k: _json_data.get(k)
+            for k in budget_params
+            if k in _json_data
+            or (k == "budget_duration" and existing_budget_id is not None and budget_duration_cleared)
+        }
+    )
 
     # Check if budget_id is explicitly provided in the data
     data_budget_id: Final[str | None] = getattr(data, "budget_id", None)

tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py

@@ -1292,6 +1292,21 @@ async def test_update_project_leaves_metadata_untouched_when_no_limit_is_sent(mo
     assert "metadata" not in _written_project_data(mock_prisma)
 
 
+@pytest.mark.asyncio
+async def test_update_project_clears_only_the_explicit_budget_cap(monkeypatch):
+    mock_prisma = _project_update_mocks(monkeypatch, {})
+    mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-clear-test"
+    mock_prisma.db.litellm_budgettable.update = mock.AsyncMock()
+
+    await _run_project_update("project-clear-test", max_budget=None)
+
+    mock_prisma.db.litellm_budgettable.update.assert_awaited_once_with(
+        where={"budget_id": "budget-clear-test"},
+        data={"max_budget": None, "updated_by": "1234"},
+    )
+    assert "max_budget" not in _written_project_data(mock_prisma)
+
+
 @pytest.mark.parametrize("entry", ["all-proxy-models", "*", "azure/*"])
 def test_enforce_project_model_quota_rejects_entries_that_expand_at_request_time(entry):
     """A quota keyed on a wildcard entry is never applied by the limiter, so it fails loudly."""

tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py

@@ -340,7 +340,8 @@ async def test_update_budget_recomputes_reset_at_when_duration_changes(
 
 
 @pytest.mark.asyncio
-async def test_update_budget_preserves_explicit_reset_at(client_and_mocks):
+@pytest.mark.parametrize("budget_duration", ["1d", None])
+async def test_update_budget_preserves_explicit_reset_at(client_and_mocks, budget_duration):
     """An explicit budget_reset_at from the caller always wins over recompute."""
     client, _, mock_table = client_and_mocks
     captured = _capture_update_data(mock_table)
@@ -350,7 +351,7 @@ async def test_update_budget_preserves_explicit_reset_at(client_and_mocks):
         "/budget/update",
         json={
             "budget_id": "budget_explicit_reset",
-            "budget_duration": "1d",
+            "budget_duration": budget_duration,
             "budget_reset_at": explicit.isoformat(),
         },
     )
@@ -377,8 +378,7 @@ async def test_update_budget_without_duration_leaves_reset_at_untouched(
 
 
 @pytest.mark.asyncio
-async def test_update_budget_duration_none_does_not_recompute(client_and_mocks):
-    """Clearing budget_duration (explicit null) must not recompute against a None duration."""
+async def test_update_budget_duration_none_clears_obsolete_reset(client_and_mocks):
     client, _, mock_table = client_and_mocks
     captured = _capture_update_data(mock_table)
 
@@ -389,7 +389,7 @@ async def test_update_budget_duration_none_does_not_recompute(client_and_mocks):
     assert resp.status_code == 200, resp.text
 
     assert "budget_duration" in captured and captured["budget_duration"] is None
-    assert "budget_reset_at" not in captured
+    assert captured["budget_reset_at"] is None
 
 
 @pytest.mark.asyncio

tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -2097,6 +2097,22 @@ def test_update_internal_user_params_reset_max_budget_with_none():
     assert non_default_values["user_id"] == "test_user"
 
 
+def test_update_internal_user_params_explicit_duration_clear_overrides_role_default(monkeypatch):
+    import litellm
+
+    monkeypatch.setattr(litellm, "internal_user_budget_duration", "30d")
+    data = UpdateUserRequest(
+        user_id="duration-clear-test",
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        budget_duration=None,
+    )
+
+    updated = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data)
+
+    assert updated["budget_duration"] is None
+    assert updated["budget_reset_at"] is None
+
+
 def test_update_internal_user_params_ignores_other_nones():
     """
     Test that other fields are still filtered out if None

ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts

@@ -10,7 +10,7 @@ export interface ProjectUpdateParams {
   description?: string;
   team_id?: string;
   models?: string[];
-  max_budget?: number;
+  max_budget?: number | null;
   blocked?: boolean;
   guardrails?: string[];
   metadata?: Record<string, unknown>;

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx

@@ -98,7 +98,11 @@ const AccessGroupBudgetModal: React.FC<AccessGroupBudgetModalProps> = ({
                 )}
               >
                 {({ id, value, onChange }) => (
-                  <BudgetDurationDropdown id={id} value={value || null} onChange={onChange} />
+                  <BudgetDurationDropdown
+                    id={id}
+                    value={value || null}
+                    onChange={(next) => onChange(next ?? undefined)}
+                  />
                 )}
               </FormField>
             </FieldGroup>

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx

@@ -102,6 +102,20 @@ describe("EditProjectModal submit payload", () => {
     });
   });
 
+  it("should send an explicit clear after blanking a saved budget", async () => {
+    const user = setup();
+    renderModal();
+
+    const budgetInput = screen.getByRole("spinbutton", { name: "Max Budget (USD)" });
+    await user.clear(budgetInput);
+    await user.tab();
+    expect(budgetInput).toHaveValue(null);
+    await save(user);
+
+    await waitFor(() => expect(mutate).toHaveBeenCalled());
+    expect(JSON.parse(JSON.stringify(variables().params))).toMatchObject({ max_budget: null });
+  });
+
   it("includes the advanced fields once Advanced Settings has been opened, even after collapsing it again", async () => {
     const user = setup();
     renderModal();

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx

@@ -86,7 +86,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit<EditProjectModalP
       : { ...values, guardrails: undefined, modelLimits: undefined, metadata: undefined };
 
     const params: ProjectUpdateParams = {
-      ...buildProjectUpdateParams(submitted),
+      ...buildProjectUpdateParams(submitted, project.litellm_budget_table?.max_budget),
       team_id: submitted.team_id,
     };
 

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx

@@ -205,8 +205,19 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr
                   type="number"
                   min={0}
                   placeholder="0.00"
-                  value={value ?? ""}
-                  onChange={(event) => onChange(toOptionalNumber(event.target.value))}
+                  value={Number.isNaN(value) ? "" : value ?? ""}
+                  onInput={(event) => {
+                    if (event.currentTarget.validity.badInput || Number.isNaN(value)) {
+                      onChange(
+                        event.currentTarget.validity.badInput
+                          ? Number.NaN
+                          : toOptionalNumber(event.currentTarget.value) ?? null,
+                      );
+                    }
+                  }}
+                  onChange={(event) =>
+                    onChange(event.target.validity.badInput ? Number.NaN : toOptionalNumber(event.target.value) ?? null)
+                  }
                 />
               </InputGroup>
             )}

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts

@@ -22,7 +22,7 @@ export const projectFormSchema = z
       .pipe(z.string({ error: "Please select a team" }).min(1, "Please select a team")),
     description: z.string().optional(),
     models: z.array(z.string()),
-    max_budget: z.number().optional(),
+    max_budget: z.number().nullish(),
     isBlocked: z.boolean(),
     guardrails: z.array(z.string()).optional(),
     modelLimits: z.array(modelLimitSchema).optional(),

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts

@@ -27,9 +27,9 @@ describe("buildProjectCreateParams", () => {
     expect(result.description).toBe("A description");
   });
 
-  it("should pass through max_budget when provided", () => {
-    const result = buildProjectCreateParams({ ...baseValues, max_budget: 50.0 });
-    expect(result.max_budget).toBe(50.0);
+  it.each([50.0, 1e308])("should preserve a finite max_budget of %s", (maxBudget) => {
+    const result = buildProjectCreateParams({ ...baseValues, max_budget: maxBudget });
+    expect(JSON.parse(JSON.stringify(result)).max_budget).toBe(maxBudget);
   });
 
   it("should build model_rpm_limit from modelLimits entries", () => {

ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts

@@ -16,6 +16,11 @@ const buildModelLimitMap = (
 const buildMetadata = (entries: ProjectFormValues["metadata"]): Record<string, string> | undefined =>
   entries && Object.fromEntries(entries.flatMap((entry) => (entry.key ? [[entry.key, entry.value] as const] : [])));
 
+const roundBudget = (value: number): number => {
+  const rounded = Math.round(value * 100) / 100;
+  return Number.isFinite(rounded) ? rounded : value;
+};
+
 const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => {
   const limitEntries = values.modelLimits ?? [];
   const modelRpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.rpm);
@@ -35,7 +40,7 @@ const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) =>
     project_alias: values.project_alias,
     description: values.description,
     models: values.models ?? [],
-    max_budget: values.max_budget === undefined ? undefined : Math.round(values.max_budget * 100) / 100,
+    max_budget: values.max_budget == null ? undefined : roundBudget(values.max_budget),
     blocked: values.isBlocked ?? false,
     ...guardrailsParam,
     ...(keep(modelRpmLimit) && { model_rpm_limit: modelRpmLimit }),
@@ -53,4 +58,7 @@ export const buildProjectCreateParams = (values: ProjectFormValues) => buildProj
  * /project/update leaves an omitted key untouched, so a limit the operator cleared has to go out as
  * an explicitly empty map. Omitting it is what silently kept a removed quota enforced.
  */
-export const buildProjectUpdateParams = (values: ProjectFormValues) => buildProjectApiParams(values, true);
+export const buildProjectUpdateParams = (values: ProjectFormValues, savedMaxBudget?: number | null) => ({
+  ...buildProjectApiParams(values, true),
+  ...(values.max_budget == null && savedMaxBudget != null ? { max_budget: null } : {}),
+});

ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx

@@ -143,7 +143,11 @@ const CreateTagModal: React.FC<CreateTagModalProps> = ({ visible, onCancel, onSu
                     )}
                   >
                     {({ id, value, onChange }) => (
-                      <BudgetDurationDropdown id={id} value={value ?? null} onChange={onChange} />
+                      <BudgetDurationDropdown
+                        id={id}
+                        value={value ?? null}
+                        onChange={(next) => onChange(next ?? undefined)}
+                      />
                     )}
                   </FormField>
                 </FieldGroup>

ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx

@@ -28,7 +28,7 @@ const tagEditShape = {
   description: z.string().optional(),
   models: z.array(z.string()).optional(),
   max_budget: z.union([z.string(), z.number()]).optional(),
-  budget_duration: z.string().optional(),
+  budget_duration: z.string().nullish(),
 };
 
 const tagEditSchema = z.object(tagEditShape);

ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx

@@ -1,4 +1,4 @@
-import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../../../../tests/test-utils";
 import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
 import { describe, expect, it, vi, beforeEach } from "vitest";
 import UserInfoView from "./user_info_view";
@@ -163,6 +163,35 @@ describe("UserInfoView add-to-team form", () => {
 
       expect(await openEditor(user)).toHaveValue(42);
     });
+
+    it("should keep Unlimited selected after saving and reopening the user", async () => {
+      const user = setup();
+      render(<UserInfoView {...budgetProps} />);
+
+      await openEditor(user);
+      await user.click(screen.getByRole("checkbox", { name: "Unlimited Budget" }));
+      await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+      await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled());
+      expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ max_budget: null });
+      await openEditor(user);
+      expect(screen.getByRole("checkbox", { name: "Unlimited Budget" })).toBeChecked();
+    });
+
+    it("should keep a cleared reset period after saving and reopening the user", async () => {
+      const user = setup();
+      render(<UserInfoView {...budgetProps} />);
+
+      await openEditor(user);
+      await user.click(screen.getByRole("combobox", { name: "Reset Budget" }));
+      await user.click(await screen.findByRole("option", { name: "n/a" }));
+      await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+      await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled());
+      expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ budget_duration: null });
+      await openEditor(user);
+      expect(screen.getByRole("combobox", { name: "Reset Budget" })).toHaveTextContent("n/a");
+    });
   });
 
   it("offers only the teams the user is not already a member of", async () => {

ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx

@@ -332,8 +332,9 @@ export default function UserInfoView({
         user_email: formValues.user_email ?? userData.user_email,
         user_alias: formValues.user_alias ?? userData.user_alias,
         models: formValues.models ?? userData.models,
-        max_budget: formValues.max_budget ?? userData.max_budget,
-        budget_duration: formValues.budget_duration ?? userData.budget_duration,
+        max_budget: formValues.max_budget === undefined ? userData.max_budget : formValues.max_budget,
+        budget_duration:
+          formValues.budget_duration === undefined ? userData.budget_duration : formValues.budget_duration,
         metadata: formValues.metadata ?? userData.metadata,
         model_max_budget: formValues.model_max_budget ?? userData.model_max_budget,
         object_permission: mcpEntitlement

ui/litellm-dashboard/src/components/Teams.tsx

@@ -807,7 +807,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
                         showNeverResets
                         placeholder={budgetDurationPlaceholder}
                         value={value}
-                        onChange={onChange}
+                        onChange={(next) => onChange(next ?? undefined)}
                       />
                     )}
                   </FormField>

ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx

@@ -14,7 +14,7 @@ const DURATION_LABELS: Record<string, string> = {
 interface BudgetDurationDropdownProps {
   id?: string;
   value?: string | null;
-  onChange?: (value: string | undefined) => void;
+  onChange?: (value: string | null) => void;
   className?: string;
   style?: React.CSSProperties;
   placeholder?: string;
@@ -31,11 +31,7 @@ const BudgetDurationDropdown: React.FC<BudgetDurationDropdownProps> = ({
   showNeverResets = false,
 }) => {
   return (
-    <Select
-      items={DURATION_LABELS}
-      value={value || null}
-      onValueChange={(next: string | null) => onChange?.(next ?? undefined)}
-    >
+    <Select items={DURATION_LABELS} value={value || null} onValueChange={onChange}>
       <SelectTrigger id={id} className={`w-full ${className}`} style={style}>
         <SelectValue placeholder={placeholder} />
       </SelectTrigger>

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -1021,7 +1021,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
                             value={control.value as string | null | undefined}
                             showNeverResets
                             placeholder="Not set"
-                            onChange={control.onChange}
+                            onChange={(next) => control.onChange(next ?? undefined)}
                           />
                         )}
                       </MountedFormField>

ui/litellm-dashboard/src/components/tag_management/types.tsx

@@ -41,7 +41,7 @@ export interface TagUpdateRequest {
   soft_budget?: number;
   tpm_limit?: number;
   rpm_limit?: number;
-  budget_duration?: string;
+  budget_duration?: string | null;
 }
 
 export interface TagDeleteRequest {

ui/litellm-dashboard/src/components/team/EditMembership.tsx

@@ -157,7 +157,7 @@ const MemberModal = <T extends BaseMember>({
               <BudgetDurationDropdown
                 id={id}
                 value={typeof value === "string" ? value : null}
-                onChange={(next) => onChange(next)}
+                onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)}
               />
             );
           default:

ui/litellm-dashboard/src/components/team/TeamInfo.tsx

@@ -1467,7 +1467,9 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                               showNeverResets
                               placeholder="Inherit team reset period"
                               value={value === null ? NEVER_RESETS_BUDGET_DURATION : value}
-                              onChange={(next) => onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next)}
+                              onChange={(next) =>
+                                onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next ?? undefined)
+                              }
                             />
                           )}
                         </FormField>