diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 5db919659ed..e19689f120a 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -66,10 +66,8 @@ class ResetBudgetJob: try: from litellm.proxy.proxy_server import spend_counter_cache - memberships = ( - await self.prisma_client.db.litellm_teammembership.find_many( - where={"budget_id": {"in": budget_ids}} - ) + memberships = await self.prisma_client.db.litellm_teammembership.find_many( + where={"budget_id": {"in": budget_ids}} ) for m in memberships: counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" @@ -574,7 +572,11 @@ class ResetBudgetJob: counter_key = None if item_type == "key" and hasattr(item, "token") and item.token is not None: counter_key = f"spend:key:{item.token}" - elif item_type == "team" and hasattr(item, "team_id") and item.team_id is not None: + elif ( + item_type == "team" + and hasattr(item, "team_id") + and item.team_id is not None + ): counter_key = f"spend:team:{item.team_id}" if counter_key is not None: @@ -637,6 +639,25 @@ class ResetBudgetJob: ) -> Optional[LiteLLM_EndUserTable]: try: enduser.spend = 0.0 + + # Invalidate the cached end-user object so auth checks pick up + # the reset spend instead of reading a stale value until TTL + # expires. Uses the same cache key format as get_end_user_object + # in litellm/proxy/auth/auth_checks.py. + if enduser.user_id is not None: + try: + from litellm.proxy.proxy_server import user_api_key_cache + + await user_api_key_cache.async_delete_cache( + key="end_user_id:{}".format(enduser.user_id) + ) + except Exception as cache_err: + verbose_proxy_logger.warning( + "Failed to invalidate cached end-user object for %s: %s. " + "Budget may be over-enforced until cache entry expires.", + enduser.user_id, + cache_err, + ) except Exception as e: verbose_proxy_logger.exception( "Error resetting budget for enduser: %s. Item: %s", e, enduser 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 f975460836a..bc351829a5f 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 @@ -252,6 +252,53 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): assert updated_budget.budget_reset_at > now +def test_reset_budget_for_enduser_invalidates_cache(): + """ + Test that _reset_budget_for_enduser invalidates the cached end-user object + so auth checks pick up the reset spend instead of a stale cached value. + """ + from unittest.mock import AsyncMock, patch + + test_enduser = type( + "LiteLLM_EndUserTable", + (), + { + "spend": 45.0, + "user_id": "cached-enduser-1", + }, + ) + + mock_cache = AsyncMock() + mock_cache.async_delete_cache = AsyncMock() + + async def _run(): + with patch( + "litellm.proxy.common_utils.reset_budget_job.user_api_key_cache", + mock_cache, + create=True, + ): + # Patch the import inside the function + with patch.dict( + "sys.modules", + { + "litellm.proxy.proxy_server": type( + "module", (), {"user_api_key_cache": mock_cache} + )() + }, + ): + result = await ResetBudgetJob._reset_budget_for_enduser( + enduser=test_enduser + ) + + assert result is not None + assert result.spend == 0.0 + mock_cache.async_delete_cache.assert_called_once_with( + key="end_user_id:cached-enduser-1" + ) + + asyncio.run(_run()) + + def test_reset_budget_all(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime now = datetime.now(timezone.utc) @@ -434,9 +481,7 @@ def test_reset_budget_for_keys_linked_to_budgets_empty( """ # Run with empty list asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets( - budgets_to_reset=[] - ) + reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[]) ) # Verify no update_many calls were made