ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Fix] Budget reset job now resets implicitly-created end users with NULL budget_id

Browse Source 
When litellm.max_end_user_budget_id is configured, implicitly-created end users
(via /chat/completions) have budget_id=NULL in the DB since the default budget
is only applied in-memory. The budget reset job filtered by budget_id, so these
users were never reset and eventually permanently blocked.

Fix: when the default budget is in the reset list, also query for and reset
end users with budget_id=NULL and spend > 0. This keeps the hot auth path
unchanged (no DB writes on every request).

Fixes #22019

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia 2026-04-04 18:10:05 -07:00 committed by shivam
parent 89365628c9
commit bb9955beca
No known key found for this signature in database
4 changed files with 247 additions and 155 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
litellm/proxy/auth/auth_checks.py
View File

@ -857,22 +857,8 @@ async def _apply_default_budget_to_end_user(
)
if default_budget is not None:
# Apply default budget to end user object (in-memory for this request)
# Apply default budget to end user object
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}"
)

47
litellm/proxy/common_utils/reset_budget_job.py
View File

@ -4,6 +4,7 @@ import time
from datetime import datetime, timezone
from typing import List, Literal, Optional, Union
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
LiteLLM_BudgetTableFull,
@ -162,16 +163,35 @@ class ResetBudgetJob:
table_name="budget",
)
budget_ids_to_reset = [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
endusers_to_reset = await self.prisma_client.get_data(
table_name="enduser",
query_type="find_all",
budget_id_list=[
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
],
budget_id_list=budget_ids_to_reset,
)
# Also reset end users with no budget_id (NULL) who use the
# default budget via litellm.max_end_user_budget_id. These
# users are enforced in-memory but never had budget_id
# persisted, so the query above misses them.
if (
litellm.max_end_user_budget_id is not None
and litellm.max_end_user_budget_id in budget_ids_to_reset
):
default_budget_endusers = (
await self._get_endusers_with_no_budget_id()
)
if default_budget_endusers:
if endusers_to_reset is None:
endusers_to_reset = default_budget_endusers
else:
endusers_to_reset.extend(default_budget_endusers)
await self.reset_budget_for_litellm_team_members(
budgets_to_reset=budgets_to_reset
)
@ -279,6 +299,23 @@ class ResetBudgetJob:
)
verbose_proxy_logger.exception("Failed to reset budget for endusers: %s", e)
async def _get_endusers_with_no_budget_id(
self,
) -> List[LiteLLM_EndUserTable]:
"""
Fetch end users that have no explicit budget_id set (NULL) and have
accumulated spend > 0. These are implicitly-created end users that
rely on the default budget (litellm.max_end_user_budget_id) applied
in-memory during auth checks.
"""
rows = await self.prisma_client.db.litellm_endusertable.find_many(
where={
"budget_id": None,
"spend": {"gt": 0},
},
)
return [LiteLLM_EndUserTable(**row.dict()) for row in rows]
async def reset_budget_for_litellm_keys(self):
"""
Resets the budget for all the litellm keys

154
tests/proxy_unit_tests/test_default_end_user_budget_simple.py
View File

@ -29,14 +29,14 @@ async def test_default_budget_applied_to_end_user_without_budget():
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,
rpm_limit=2,
tpm_limit=10,
)
# Mock end user in DB without budget
mock_end_user_data = {
"user_id": end_user_id,
@ -47,12 +47,11 @@ async def test_default_budget_applied_to_end_user_without_budget():
"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_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: default_budget.dict())
)
@ -76,12 +75,6 @@ async def test_default_budget_applied_to_end_user_without_budget():
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
@ -95,13 +88,13 @@ async def test_explicit_budget_not_overridden_by_default():
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,
rpm_limit=50,
)
# Mock end user with explicit budget
mock_end_user_data = {
"user_id": end_user_id,
@ -112,29 +105,29 @@ async def test_explicit_budget_not_overridden_by_default():
"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_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 explicit budget is kept (not replaced with default)
assert result is not None
assert result.litellm_budget_table.budget_id == explicit_budget_id
assert result.litellm_budget_table.max_budget == 100.0
assert result.litellm_budget_table.rpm_limit == 50
litellm.max_end_user_budget_id = None
@ -147,13 +140,13 @@ async def test_budget_enforcement_blocks_over_budget_users():
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,
rpm_limit=2,
)
# Mock end user who has already spent more than budget
mock_end_user_data = {
"user_id": end_user_id,
@ -164,12 +157,11 @@ async def test_budget_enforcement_blocks_over_budget_users():
"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_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: default_budget.dict())
)
@ -186,116 +178,9 @@ async def test_budget_enforcement_blocks_over_budget_users():
user_api_key_cache=mock_cache,
route="/chat/completions",
)
assert "ExceededBudget" in str(exc_info.value)
assert end_user_id in str(exc_info.value)
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
@ -308,7 +193,7 @@ async def test_system_works_without_default_budget_configured():
"""
end_user_id = f"test_user_{uuid.uuid4().hex}"
litellm.max_end_user_budget_id = None # Not configured
# Mock end user without budget
mock_end_user_data = {
"user_id": end_user_id,
@ -319,25 +204,24 @@ async def test_system_works_without_default_budget_configured():
"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_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 work fine, just without budget limits
assert result is not None
assert result.user_id == end_user_id
assert result.litellm_budget_table is None # No budget applied

185
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View File

@ -36,10 +36,24 @@ class MockLiteLLMVerificationToken:
return {"count": 1}
class MockLiteLLMEndUserTable:
def __init__(self):
self.find_many_calls: List[Dict[str, Any]] = []
self._find_many_results: List[Any] = []
def set_find_many_results(self, results: List[Any]):
self._find_many_results = results
async def find_many(self, where: Dict[str, Any]) -> List[Any]:
self.find_many_calls.append({"where": where})
return self._find_many_results
class MockDB:
def __init__(self):
self.litellm_teammembership = MockLiteLLMTeamMembership()
self.litellm_verificationtoken = MockLiteLLMVerificationToken()
self.litellm_endusertable = MockLiteLLMEndUserTable()
class MockPrismaClient:
@ -613,3 +627,174 @@ def test_budget_table_reset_also_resets_linked_keys(
)
assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]}
assert calls[0]["data"]["spend"] == 0
def test_reset_budget_resets_endusers_with_null_budget_id(
reset_budget_job, mock_prisma_client
):
"""
When litellm.max_end_user_budget_id is configured and that budget is
being reset, end users with budget_id=NULL should also have their spend
reset. These users were implicitly created and have no budget_id persisted,
but are enforced against the default budget in-memory.
"""
import litellm
now = datetime.now(timezone.utc)
default_budget_id = "default-enduser-budget"
litellm.max_end_user_budget_id = default_budget_id
# Budget that is due for reset — matches the default end user budget
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 50.0,
"budget_duration": "1d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": default_budget_id,
"created_at": now - timedelta(days=1),
},
)
# End user WITH explicit budget_id (found by the normal budget_id_list query)
enduser_with_budget = type(
"LiteLLM_EndUserTable",
(),
{
"spend": 30.0,
"litellm_budget_table": test_budget,
"user_id": "enduser-explicit",
},
)
# End user WITHOUT budget_id (NULL) — should also be reset
enduser_no_budget_row = type(
"EndUserRow",
(),
{
"spend": 25.0,
"user_id": "enduser-implicit",
"budget_id": None,
"alias": None,
"allowed_model_region": None,
"default_model": None,
"blocked": False,
"object_permission_id": None,
"object_permission": None,
"litellm_budget_table": None,
"dict": lambda self=None: {
"spend": 25.0,
"user_id": "enduser-implicit",
"blocked": False,
"alias": None,
"allowed_model_region": None,
"default_model": None,
"litellm_budget_table": None,
"object_permission_id": None,
"object_permission": None,
},
},
)
mock_prisma_client.data["budget"] = [test_budget]
mock_prisma_client.data["enduser"] = [enduser_with_budget]
# Set up the DB mock for NULL-budget-id end users
mock_prisma_client.db.litellm_endusertable.set_find_many_results(
[enduser_no_budget_row]
)
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
# Both end users should have been reset
updated = mock_prisma_client.updated_data["enduser"]
assert len(updated) == 2, (
f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}"
)
user_ids = {u.user_id for u in updated}
assert "enduser-explicit" in user_ids
assert "enduser-implicit" in user_ids
for u in updated:
assert u.spend == 0.0, f"Expected spend=0 for {u.user_id}, got {u.spend}"
# Verify find_many was called to fetch NULL-budget-id end users
find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert len(find_many_calls) == 1
assert find_many_calls[0]["where"] == {"budget_id": None, "spend": {"gt": 0}}
litellm.max_end_user_budget_id = None
def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured(
reset_budget_job, mock_prisma_client
):
"""
When litellm.max_end_user_budget_id is NOT configured, end users with
budget_id=NULL should NOT be fetched or reset.
"""
import litellm
now = datetime.now(timezone.utc)
litellm.max_end_user_budget_id = None
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 50.0,
"budget_duration": "1d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": "some-budget",
"created_at": now - timedelta(days=1),
},
)
mock_prisma_client.data["budget"] = [test_budget]
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
# Should NOT have queried for NULL-budget-id end users
find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert len(find_many_calls) == 0
litellm.max_end_user_budget_id = None
def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_list(
reset_budget_job, mock_prisma_client
):
"""
When litellm.max_end_user_budget_id IS configured but the corresponding
budget is NOT in the budgets-to-reset list (not yet expired), end users
with budget_id=NULL should NOT be reset.
"""
import litellm
now = datetime.now(timezone.utc)
litellm.max_end_user_budget_id = "default-budget-not-expired"
# A different budget that IS expiring (not the default one)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 50.0,
"budget_duration": "1d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": "other-budget",
"created_at": now - timedelta(days=1),
},
)
mock_prisma_client.data["budget"] = [test_budget]
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
# Should NOT have queried for NULL-budget-id end users
find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert len(find_many_calls) == 0
litellm.max_end_user_budget_id = None