From 8604d4518dcfba65cef504aa6c2f46a727e1ad20 Mon Sep 17 00:00:00 2001
From: Adam Dickinson <adickinson@demandbase.com>
Date: Wed, 8 Apr 2026 19:56:29 -0500
Subject: [PATCH] fix: persist budget_id to DB for auto-created end users using
 max_end_user_budget_id
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When `max_end_user_budget_id` is configured, end users created via the
`x-litellm-end-user-id` header have their default budget applied only
in-memory at auth time. The `budget_id` is never written to the
`LiteLLM_EndUserTable` row, so the budget reset job (which queries
`WHERE budget_id IN (...)`) never finds these users — their spend
accumulates forever and never resets.

This commit:
1. Includes `budget_id` in the upsert create clause in
   `update_end_user_spend()` when the default budget is configured,
   so newly created end users get `budget_id` persisted from the start.
2. Adds a lazy backfill in `_apply_default_budget_to_end_user()` that
   writes `budget_id` to the DB for existing users where it is NULL,
   so pre-existing users self-heal on their next request.
3. Adds the `budget_id` field to the `LiteLLM_EndUserTable` Pydantic
   model to match the Prisma schema.

Fixes #25386
---
 litellm/proxy/_types.py           |  1 +
 litellm/proxy/auth/auth_checks.py | 23 ++++++++++++++++++++---
 litellm/proxy/utils.py            | 13 ++++++++-----
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index cf99c5cd9fa..590a81e7338 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2897,6 +2897,7 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase):
     spend: float = 0.0
     allowed_model_region: Optional[AllowedModelRegion] = None
     default_model: Optional[str] = None
+    budget_id: Optional[str] = None
     litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
     object_permission_id: Optional[str] = None
     object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 68bde8434a6..98ed24773bf 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -918,9 +918,26 @@ async def _apply_default_budget_to_end_user(
     if default_budget is not None:
         # Apply default budget to end user object
         end_user_obj.litellm_budget_table = default_budget
-        verbose_proxy_logger.debug(
-            f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}"
-        )
+
+        # Backfill budget_id to DB for existing users that were created without it
+        if end_user_obj.budget_id is None:
+            try:
+                await prisma_client.db.litellm_endusertable.update(
+                    where={"user_id": end_user_obj.user_id},
+                    data={"budget_id": litellm.max_end_user_budget_id},
+                )
+                end_user_obj.budget_id = litellm.max_end_user_budget_id
+                verbose_proxy_logger.debug(
+                    f"Persisted default budget_id {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}"
+                )
+            except Exception:
+                verbose_proxy_logger.debug(
+                    f"Failed to persist default budget_id for end user {end_user_obj.user_id}"
+                )
+        else:
+            verbose_proxy_logger.debug(
+                f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}"
+            )
 
     return end_user_obj
 
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index ec98cfd4d1e..32730877918 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -4679,14 +4679,17 @@ class ProxyUpdateSpend:
                         ) in end_user_list_transactions.items():
                             if litellm.max_end_user_budget is not None:
                                 pass
+                            _create_data: dict = {
+                                "user_id": end_user_id,
+                                "spend": response_cost,
+                                "blocked": False,
+                            }
+                            if litellm.max_end_user_budget_id is not None:
+                                _create_data["budget_id"] = litellm.max_end_user_budget_id
                             batcher.litellm_endusertable.upsert(
                                 where={"user_id": end_user_id},
                                 data={
-                                    "create": {
-                                        "user_id": end_user_id,
-                                        "spend": response_cost,
-                                        "blocked": False,
-                                    },
+                                    "create": _create_data,
                                     "update": {"spend": {"increment": response_cost}},
                                 },
                             )