diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dd81afe78c..56958a88f6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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}" ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 16243038b7..d935c50ffd 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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 diff --git a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py index 8c85a79cbf..970a7ab471 100644 --- a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py +++ b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py @@ -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 - 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 1f2d4f4905..db531bb3bb 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 @@ -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