From 04a2407244bc785d47f173c7efee1283122ad3fa Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 23:03:14 +0000 Subject: [PATCH] fix(organization): clear org budget limits when PATCH /organization/update sends null A sent null for tpm_limit, rpm_limit, max_budget and the other budget fields was dropped by a 'v is not None' filter, so update_budget was never called and the request returned 200 without changing the budget row. Presence is now read from model_fields_set (merge-patch semantics, matching /v2/organization) and the nested litellm_budget_table payload no longer drops nulls either Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../organization_endpoints.py | 5 +- .../test_organization_endpoints.py | 58 +++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 5e38a016099..d32b842df58 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -312,7 +312,7 @@ def handle_nested_budget_structure_in_organization_update_request( # Extract valid budget fields and merge into top level budget_fields: Final = LiteLLM_BudgetTable.model_fields.keys() for key, value in budget_data.items(): - if key in budget_fields and value is not None: + if key in budget_fields: transformed_data[key] = value return transformed_data @@ -708,9 +708,8 @@ async def update_organization( existing_organization_row=existing_organization_row, ) - # Handle budget updates if budget fields are provided budget_fields: Final = { - k: v for k, v in data.model_dump().items() if k in LiteLLM_BudgetTable.model_fields and v is not None + k: v for k, v in data.model_dump().items() if k in _BUDGET_SETTABLE_FIELDS and k in data.model_fields_set } if budget_fields and existing_organization_row.budget_id: diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index e2d89a660c2..5c2a0bdde3d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -963,6 +963,64 @@ async def test_v2_serializes_model_max_budget_on_budget_write(monkeypatch): assert json.loads(written) == {"gpt-4o": {"max_budget": 10}} +async def _run_legacy_update_organization(monkeypatch, *, body: dict, existing_budget_id: str): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import update_organization + from litellm.proxy.utils import jsonify_object + + mock_prisma_client = AsyncMock() + mock_prisma_client.jsonify_object = jsonify_object + + existing_org = MagicMock() + existing_org.budget_id = existing_budget_id + existing_org.metadata = {} + mock_prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(return_value=existing_org) + mock_prisma_client.db.litellm_organizationtable.update = AsyncMock(return_value=MagicMock()) + mock_prisma_client.db.litellm_budgettable.update = AsyncMock() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_verify_org_access", AsyncMock()) + + request = MagicMock() + request.json = AsyncMock(return_value=body) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + await update_organization(request=request, user_api_key_dict=auth) + return mock_prisma_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + {"organization_id": "org-1", "tpm_limit": None}, + {"organization_id": "org-1", "litellm_budget_table": {"tpm_limit": None}}, + ], +) +async def test_legacy_update_clears_tpm_limit_when_sent_null(monkeypatch, body): + """PATCH /organization/update with tpm_limit: null writes None to the budget row instead of dropping it.""" + prisma = await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1") + + budget_write = prisma.db.litellm_budgettable.update.await_args + assert budget_write.kwargs["where"] == {"budget_id": "budget-1"} + assert budget_write.kwargs["data"]["tpm_limit"] is None + assert "rpm_limit" not in budget_write.kwargs["data"] + assert "tpm_limit" not in prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_legacy_update_without_budget_fields_skips_budget_write(monkeypatch): + """Omitted budget fields are left untouched: renaming the org must not write the budget row.""" + prisma = await _run_legacy_update_organization( + monkeypatch, + body={"organization_id": "org-1", "organization_alias": "renamed"}, + existing_budget_id="budget-1", + ) + + prisma.db.litellm_budgettable.update.assert_not_awaited() + assert prisma.db.litellm_organizationtable.update.await_args.kwargs["data"]["organization_alias"] == "renamed" + + def test_build_budget_write_data_recomputes_reset_at_on_duration(): """A sent budget_duration recomputes budget_reset_at so the reset window follows the new duration.""" from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data