From ffe56cf5a2bf93ffd5cbd60e6086810d9233cf3b Mon Sep 17 00:00:00 2001 From: Joseph Yeung Date: Mon, 20 Jul 2026 19:55:38 -0400 Subject: [PATCH] fix(budget): reset users/teams with NULL budget_reset_at (#33623) Internal users seeded from default_internal_user_params (SSO/JWT first-login upsert, or /user/new without an explicit budget_reset_at) get budget_duration set but budget_reset_at = NULL. The ResetBudgetJob user/team queries filter on {"budget_reset_at": {"lt": now}}, which never matches NULL, so these rows are never reset: their spend accumulates for the lifetime of the row and silently exceeds max_budget with no periodic reset. The budget-table query already handles this by OR-ing in a {budget_reset_at IS NULL AND budget_duration IS NOT NULL} branch. Apply the same pattern to the user and team reset queries in PrismaClient.get_data. Adds a regression test asserting both the user and team reset queries select NULL-budget_reset_at rows that have a budget_duration. Co-authored-by: yuneng-jiang Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 --- litellm/proxy/utils.py | 32 ++++++++- .../common_utils/test_reset_budget_job.py | 68 +++++++++++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 844d3c2ed26..43921a847a9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3320,7 +3320,24 @@ class PrismaClient: elif query_type == "find_all" and reset_at is not None: response = await UserRepository(self).table.find_many( where={ # type: ignore - "budget_reset_at": {"lt": reset_at}, + # A user seeded from default_internal_user_params + # (or created via /user/new without an explicit + # budget_reset_at) has budget_duration set but + # budget_reset_at = NULL. `{"lt": reset_at}` never + # matches NULL, so such users would never be reset + # and their spend would accumulate for the lifetime + # of the row, silently exceeding max_budget. Treat a + # NULL budget_reset_at with a non-NULL budget_duration + # as due, matching the budget-table query below. + "OR": [ + { + "AND": [ + {"budget_reset_at": None}, + {"NOT": {"budget_duration": None}}, + ] + }, + {"budget_reset_at": {"lt": reset_at}}, + ], } ) elif query_type == "find_all" and user_id_list is not None: @@ -3406,7 +3423,18 @@ class PrismaClient: elif query_type == "find_all" and reset_at is not None: response = await TeamRepository(self).table.find_many( where={ # type: ignore - "budget_reset_at": {"lt": reset_at}, + # Same NULL budget_reset_at gap as the user query + # above: a team with a budget_duration but no + # initialized budget_reset_at would never be reset. + "OR": [ + { + "AND": [ + {"budget_reset_at": None}, + {"NOT": {"budget_duration": None}}, + ] + }, + {"budget_reset_at": {"lt": reset_at}}, + ], } ) elif query_type == "find_all" and user_id is not None: diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 0b683745369..5e348b1bb7e 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1803,3 +1803,71 @@ def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) prisma_client.db.litellm_tagtable.update_many.assert_awaited_once() + + +def _extract_reset_where(find_many_mock): + """Return the ``where`` dict passed to a mocked repository ``find_many``.""" + assert find_many_mock.await_count == 1 + _, kwargs = find_many_mock.await_args + return kwargs["where"] + + +def _asserts_null_reset_is_due(where): + """A budget-reset ``find_many`` filter must select rows whose + ``budget_reset_at`` is NULL but which have a ``budget_duration`` set, in + addition to rows whose ``budget_reset_at`` is already in the past. + + Regression guard: a user/team seeded from ``default_internal_user_params`` + (or created via ``/user/new`` without an explicit ``budget_reset_at``) has + ``budget_duration`` set but ``budget_reset_at = NULL``. A plain + ``{"budget_reset_at": {"lt": now}}`` filter never matches NULL, so such rows + would never be reset and their spend would accumulate for the lifetime of + the row, silently exceeding ``max_budget``. + """ + branches = where.get("OR") + assert isinstance(branches, list), f"expected an OR filter, got {where!r}" + + has_null_branch = any( + b.get("AND") + == [ + {"budget_reset_at": None}, + {"NOT": {"budget_duration": None}}, + ] + for b in branches + if isinstance(b, dict) + ) + has_expired_branch = any( + isinstance(b, dict) + and "budget_reset_at" in b + and b["budget_reset_at"] is not None + for b in branches + ) + assert has_null_branch, f"missing NULL-reset_at branch in {where!r}" + assert has_expired_branch, f"missing expired-reset_at branch in {where!r}" + + +@pytest.mark.parametrize("table_name", ["user", "team"]) +def test_get_data_reset_query_selects_null_budget_reset_at(table_name): + """``PrismaClient.get_data(..., reset_at=...)`` for the user and team tables + must select rows with a NULL ``budget_reset_at`` (and a non-NULL + ``budget_duration``), matching the budget-table query. Without this, users + auto-created from ``default_internal_user_params`` are never reset.""" + from litellm.proxy.utils import PrismaClient + + # Build a PrismaClient without running its heavy __init__; only .db is used. + client = PrismaClient.__new__(PrismaClient) + client.db = MagicMock() + + find_many = AsyncMock(return_value=[]) + table_attr = { + "user": "litellm_usertable", + "team": "litellm_teamtable", + }[table_name] + setattr(getattr(client.db, table_attr), "find_many", find_many) + + now = datetime.now(timezone.utc) + asyncio.run( + client.get_data(table_name=table_name, query_type="find_all", reset_at=now) + ) + + _asserts_null_reset_is_due(_extract_reset_where(find_many))