[Fix] Persist default end-user budget_id to DB so budget reset job picks up implicitly created users

Previously, _apply_default_budget_to_end_user() only set the budget in-memory,
leaving budget_id NULL in the database. This caused the budget reset job to skip
these users since it filters by budget_id. Now the function also persists
budget_id via a Prisma update call (non-fatal on failure).

Fixes #22019

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

litellm/proxy/auth/auth_checks.py

@@ -916,8 +916,22 @@ async def _apply_default_budget_to_end_user(
     )
 
     if default_budget is not None:
-        # Apply default budget to end user object
+        # Apply default budget to end user object (in-memory for this request)
         end_user_obj.litellm_budget_table = default_budget
+
+        # Persist budget_id to DB so the budget reset job can find this user
+        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},
+            )
+        except Exception as e:
+            verbose_proxy_logger.warning(
+                "Failed to persist default budget_id for end user %s: %s",
+                end_user_obj.user_id,
+                e,
+            )
+
         verbose_proxy_logger.debug(
             f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}"
         )

tests/proxy_unit_tests/test_default_end_user_budget_simple.py

@@ -52,21 +52,22 @@ async def test_default_budget_applied_to_end_user_without_budget():
     mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
         return_value=MagicMock(dict=lambda: mock_end_user_data)
     )
+    mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None)
     mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
         return_value=MagicMock(dict=lambda: default_budget.dict())
     )
-    
+
     mock_cache = AsyncMock(spec=DualCache)
     mock_cache.async_get_cache = AsyncMock(return_value=None)
     mock_cache.async_set_cache = AsyncMock()
-    
+
     result = await get_end_user_object(
         end_user_id=end_user_id,
         prisma_client=mock_prisma_client,
         user_api_key_cache=mock_cache,
         route="/chat/completions",
     )
-    
+
     # Verify default budget was applied
     assert result is not None
     assert result.litellm_budget_table is not None
@@ -74,7 +75,13 @@ async def test_default_budget_applied_to_end_user_without_budget():
     assert result.litellm_budget_table.max_budget == 10.0
     assert result.litellm_budget_table.rpm_limit == 2
     assert result.litellm_budget_table.tpm_limit == 10
-    
+
+    # Verify budget_id was persisted to DB
+    mock_prisma_client.db.litellm_endusertable.update.assert_called_once_with(
+        where={"user_id": end_user_id},
+        data={"budget_id": default_budget_id},
+    )
+
     litellm.max_end_user_budget_id = None
 
 
@@ -162,14 +169,15 @@ async def test_budget_enforcement_blocks_over_budget_users():
     mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
         return_value=MagicMock(dict=lambda: mock_end_user_data)
     )
+    mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None)
     mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
         return_value=MagicMock(dict=lambda: default_budget.dict())
     )
-    
+
     mock_cache = AsyncMock(spec=DualCache)
     mock_cache.async_get_cache = AsyncMock(return_value=None)
     mock_cache.async_set_cache = AsyncMock()
-    
+
     # Should raise BudgetExceededError
     with pytest.raises(litellm.BudgetExceededError) as exc_info:
         await get_end_user_object(
@@ -185,6 +193,113 @@ async def test_budget_enforcement_blocks_over_budget_users():
     litellm.max_end_user_budget_id = None
 
 
+@pytest.mark.asyncio
+async def test_default_budget_db_persist_failure_is_nonfatal():
+    """
+    If the DB update to persist budget_id fails, the budget should still be
+    applied in-memory for the current request (non-fatal warning).
+    """
+    end_user_id = f"test_user_{uuid.uuid4().hex}"
+    default_budget_id = str(uuid.uuid4())
+    litellm.max_end_user_budget_id = default_budget_id
+
+    default_budget = LiteLLM_BudgetTable(
+        budget_id=default_budget_id,
+        max_budget=10.0,
+    )
+
+    mock_end_user_data = {
+        "user_id": end_user_id,
+        "spend": 1.0,
+        "litellm_budget_table": None,
+        "alias": None,
+        "allowed_model_region": None,
+        "default_model": None,
+        "blocked": False,
+    }
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
+        return_value=MagicMock(dict=lambda: mock_end_user_data)
+    )
+    # Simulate DB update failure
+    mock_prisma_client.db.litellm_endusertable.update = AsyncMock(
+        side_effect=Exception("DB connection lost")
+    )
+    mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
+        return_value=MagicMock(dict=lambda: default_budget.model_dump())
+    )
+
+    mock_cache = AsyncMock(spec=DualCache)
+    mock_cache.async_get_cache = AsyncMock(return_value=None)
+    mock_cache.async_set_cache = AsyncMock()
+
+    result = await get_end_user_object(
+        end_user_id=end_user_id,
+        prisma_client=mock_prisma_client,
+        user_api_key_cache=mock_cache,
+        route="/chat/completions",
+    )
+
+    # Budget should still be applied in-memory despite DB failure
+    assert result is not None
+    assert result.litellm_budget_table is not None
+    assert result.litellm_budget_table.budget_id == default_budget_id
+
+    litellm.max_end_user_budget_id = None
+
+
+@pytest.mark.asyncio
+async def test_explicit_budget_skips_db_update():
+    """
+    End users with an explicit budget should NOT trigger a DB update
+    for the default budget_id.
+    """
+    end_user_id = f"test_user_{uuid.uuid4().hex}"
+    explicit_budget_id = str(uuid.uuid4())
+    default_budget_id = str(uuid.uuid4())
+    litellm.max_end_user_budget_id = default_budget_id
+
+    explicit_budget = LiteLLM_BudgetTable(
+        budget_id=explicit_budget_id,
+        max_budget=100.0,
+    )
+
+    mock_end_user_data = {
+        "user_id": end_user_id,
+        "spend": 10.0,
+        "litellm_budget_table": explicit_budget.model_dump(),
+        "alias": None,
+        "allowed_model_region": None,
+        "default_model": None,
+        "blocked": False,
+    }
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
+        return_value=MagicMock(dict=lambda: mock_end_user_data)
+    )
+    mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None)
+
+    mock_cache = AsyncMock(spec=DualCache)
+    mock_cache.async_get_cache = AsyncMock(return_value=None)
+    mock_cache.async_set_cache = AsyncMock()
+
+    result = await get_end_user_object(
+        end_user_id=end_user_id,
+        prisma_client=mock_prisma_client,
+        user_api_key_cache=mock_cache,
+        route="/chat/completions",
+    )
+
+    # Should keep explicit budget, NOT call update
+    assert result is not None
+    assert result.litellm_budget_table.budget_id == explicit_budget_id
+    mock_prisma_client.db.litellm_endusertable.update.assert_not_called()
+
+    litellm.max_end_user_budget_id = None
+
+
 @pytest.mark.asyncio
 async def test_system_works_without_default_budget_configured():
     """