From 6021f31ebc8b1138b0bd032da59824833e2842a5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 5 Dec 2025 11:45:23 -0800 Subject: [PATCH] Fix: Allow null max_budget in budget update endpoint (#17545) Co-authored-by: Cursor Agent Co-authored-by: ishaan --- .../budget_management_endpoints.py | 2 +- .../test_budget_endpoints.py | 33 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 804fe274cc9..2d86f74a41c 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -110,7 +110,7 @@ async def update_budget( response = await prisma_client.db.litellm_budgettable.update( where={"budget_id": budget_obj.budget_id}, data={ - **budget_obj.model_dump(exclude_none=True), # type: ignore + **budget_obj.model_dump(exclude_unset=True), # type: ignore "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, }, # type: ignore ) 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 5dab71a1679..b4dcc33c747 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -130,3 +130,36 @@ async def test_update_budget_db_not_connected(client_and_mocks, monkeypatch): assert resp.status_code == 500 detail = resp.json()["detail"] assert detail["error"] == CommonProxyErrors.db_not_connected_error.value + + +@pytest.mark.asyncio +async def test_update_budget_allows_null_max_budget(client_and_mocks): + """ + Test that /budget/update allows setting max_budget to null. + + Previously, using exclude_none=True would drop null values, + making it impossible to remove a budget limit. With exclude_unset=True, + explicitly setting max_budget to null should include it in the update. + """ + client, _, mock_table = client_and_mocks + + captured_data = {} + + 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_789", + "max_budget": None, # Explicitly setting to null to remove budget limit + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 200, resp.text + + # Verify that max_budget=None was included in the update data + assert "max_budget" in captured_data, "max_budget should be included when explicitly set to null" + assert captured_data["max_budget"] is None, "max_budget should be None" + + mock_table.update.assert_awaited_once()