fix(proxy): clear model budget limits on explicit null

litellm/proxy/management_endpoints/budget_management_endpoints.py

@@ -141,6 +141,7 @@ async def update_budget(
     - rpm_limit: Optional[int] - The requests per minute limit for the budget.
     - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
     - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
+      Send null or {} to clear all per-model limits. Omit this field to keep them unchanged.
     - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
     """
     from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
@@ -189,9 +190,13 @@ async def update_budget(
         else {}
     )
 
+    budget_patch: Final[Mapping[str, object]] = budget_obj.model_dump(exclude_unset=True)
     budget_obj_jsonified: Final[Mapping[str, object]] = jsonify_object(
         {
-            **budget_obj.model_dump(exclude_unset=True),
+            **{
+                field: {} if field == "model_max_budget" and value is None else value
+                for field, value in budget_patch.items()
+            },
             **recomputed_reset_at,
             "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
         }

tests/integration/management/test_budget_updates.py

@@ -1,3 +1,4 @@
+import json
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from typing import Final
@@ -6,7 +7,7 @@ import pytest
 from pydantic import JsonValue, TypeAdapter
 
 from tests.integration._support.client import Gateway, string_value
-from tests.integration._support.database import read_rows
+from tests.integration._support.database import read_rows, write_rows
 from tests.integration._support.process import owned_proxy
 
 pytestmark: Final = pytest.mark.shares_worker_state
@@ -213,3 +214,38 @@ def test_model_max_budget_update_is_refused_by_the_handler_without_a_license(gat
             served: Final = unlicensed.request("POST", "/budget/update", {"budget_id": budget_id, "max_budget": 20.0})
             assert served.status_code == 200, served.text
             assert _budget_row(budget_id) == {**before, "max_budget": 20.0}
+
+
+@pytest.mark.parametrize("model_max_budget", [None, {}], ids=["null", "empty-object"])
+def test_clearing_model_budgets_preserves_other_limits(
+    gateway: Gateway, model_max_budget: dict[str, JsonValue] | None
+) -> None:
+    stored_model_budgets: Final = {
+        "qa-budget-a": {"max_budget": 1.0, "budget_duration": "1d"},
+        "qa-budget-b": {"max_budget": 2.0, "budget_duration": "7d"},
+    }
+    with gateway.scenario() as scenario:
+        budget_id: Final = scenario.budget(
+            max_budget=20.0,
+            soft_budget=5.0,
+            tpm_limit=1000,
+            rpm_limit=60,
+            budget_duration="30d",
+        )
+        write_rows(
+            'UPDATE "LiteLLM_BudgetTable" SET model_max_budget = %s::jsonb WHERE budget_id = %s',
+            (json.dumps(stored_model_budgets), budget_id),
+        )
+        before_row: Final = _budget_row(budget_id)
+        before_info: Final = _budget_info(gateway, budget_id)
+        assert before_row["model_max_budget"] == stored_model_budgets
+
+        response: Final = gateway.request(
+            "POST", "/budget/update", {"budget_id": budget_id, "model_max_budget": model_max_budget}
+        )
+
+        assert response.status_code == 200, response.text
+        assert _budget_row(budget_id) == {**before_row, "model_max_budget": {}}
+        info: Final = _budget_info(gateway, budget_id)
+        assert info == {**before_info, "model_max_budget": {}, "updated_at": info["updated_at"]}
+        assert response.json() == info, response.text

tests/unit/proxy/management_endpoints/test_budget_endpoints.py

@@ -477,3 +477,33 @@ async def test_update_budget_serializes_model_max_budget_for_prisma(
         "gpt4o": {"max_budget": 5.0, "budget_duration": "1d"},
         "glm-5.2": {"max_budget": 7.5, "budget_duration": "30d"},
     }
+
+
+@pytest.mark.parametrize(
+    ("model_patch", "expected_model_patch"),
+    [
+        pytest.param({}, {}, id="omitted"),
+        pytest.param({"model_max_budget": None}, {"model_max_budget": "{}"}, id="null"),
+        pytest.param({"model_max_budget": {}}, {"model_max_budget": "{}"}, id="empty-object"),
+    ],
+)
+def test_update_budget_preserves_omitted_model_budgets_and_clears_explicit_empty_values(
+    client_and_mocks: tuple[TestClient, MagicMock, MagicMock],
+    model_patch: dict[str, object],
+    expected_model_patch: dict[str, object],
+) -> None:
+    client, _, mock_table = client_and_mocks
+    captured: Final = _capture_update_data(mock_table)
+
+    response: Final = client.post(
+        "/budget/update",
+        json={"budget_id": "budget_model_clear", "soft_budget": None, **model_patch},
+    )
+
+    assert response.status_code == 200, response.text
+    assert captured == {
+        "budget_id": "budget_model_clear",
+        "soft_budget": None,
+        "updated_by": "test_user",
+        **expected_model_patch,
+    }

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -2099,6 +2099,7 @@ export interface paths {
          *     - rpm_limit: Optional[int] - The requests per minute limit for the budget.
          *     - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
          *     - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
+         *       Send null or {} to clear all per-model limits. Omit this field to keep them unchanged.
          *     - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
          */
         post: operations["update_budget_budget_update_post"];