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 <yuneng@berri.ai>
Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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:

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))