scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

fix(proxy): apply an explicitly sent zero customer budget instead of ignoring it

Browse source
This commit is contained in:
lewis 2026-10-16 22:44:03 -04:00
parent 363ca735ec
commit 8408cd94b5
2 changed files with 9 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -92,7 +92,7 @@ _CUSTOMER_UPDATE_LOCK_SQL: Final = 'SELECT "user_id" FROM "LiteLLM_EndUserTable"
def _should_update_field(field: str, value: object, sent_fields: AbstractSet[str]) -> bool:
if value is None:
return False
if field in sent_fields and (isinstance(value, bool) or field in _CLEARABLE_LIST_FIELDS):
if field in sent_fields and (isinstance(value, (int, float)) or field in _CLEARABLE_LIST_FIELDS):
return True
if isinstance(value, (list, dict)) and not value:
return False
@ -686,7 +686,7 @@ async def _update_single_customer(
end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump())
budget_error: Final = (
await _customer_budget_update_error(tx, end_user_table_data_typed.budget_id, data.user_id)
if "max_budget" in budget_table_data and end_user_table_data_typed.budget_id is not None
if budget_table_data and end_user_table_data_typed.budget_id is not None
else None
)
if budget_error is not None:

11
tests/unit/proxy/management_endpoints/test_customer_endpoints.py
View file

@ -73,7 +73,8 @@ client = TestClient(app)
("object_permission", {"mcp_servers": ["s1"]}, frozenset({"object_permission"}), True),
("metadata", ["m1"], frozenset(), True),
("models", ["m1"], frozenset(), True),
("max_budget", 0, frozenset({"max_budget"}), False),
("max_budget", 0, frozenset({"max_budget"}), True),
("max_budget", 0, frozenset(), False),
("max_budget", 5.0, frozenset(), True),
("alias", "a", frozenset(), True),
),
@ -89,6 +90,7 @@ client = TestClient(app)
"nonempty-list-other-field",
"nonempty-models",
"zero-budget",
"omitted-zero-budget",
"nonzero-budget",
"alias",
),
@ -2088,11 +2090,12 @@ def test_customer_budget_update_rejects_shared_references_without_partial_writes
assert database.locked_budget_ids == {budget.budget_id}
@pytest.mark.parametrize("new_max_budget", (100, 0))
@pytest.mark.parametrize(
"route", ("/customer/update", "/end_user/update", "/customer/bulk_update", "/end_user/bulk_update")
)
def test_customer_budget_update_allows_own_budget_with_unrelated_shared_references(
mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str
mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str, new_max_budget: int
) -> None:
own_budget: Final = LiteLLM_BudgetTable(budget_id="own-budget", max_budget=50)
other_budget: Final = LiteLLM_BudgetTable(budget_id="other-budget", max_budget=75)
@ -2112,13 +2115,13 @@ def test_customer_budget_update_allows_own_budget_with_unrelated_shared_referenc
(own_customer, other), budgets=(own_budget, other_budget), budget_references=references
)
database.attach(mock_prisma_client)
update: Final = {"user_id": own_customer.user_id, "max_budget": 100}
update: Final = {"user_id": own_customer.user_id, "max_budget": new_max_budget}
payload: Final = {"customers": [update]} if route.endswith("bulk_update") else update
with patch("litellm.max_end_user_budget_id", "other-budget"):
response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"})
assert response.status_code == 200, response.text
assert database.budgets[own_budget.budget_id].max_budget == 100
assert database.budgets[own_budget.budget_id].max_budget == new_max_budget
assert database.budgets[other_budget.budget_id] == other_budget
assert database.customers[own_customer.user_id].budget_id == own_budget.budget_id
assert database.locked_budget_ids == {own_budget.budget_id}