From d0383412e8a48f7fe0296aa781943ba03200b440 Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Sat, 31 Jan 2026 21:31:44 +0000 Subject: [PATCH] fix: make cache updates synchronous for budget enforcement The budget enforcement was failing in tests because cache updates were fire-and-forget (asyncio.create_task), causing race conditions where subsequent requests would read stale spend data. Changes: 1. proxy_track_cost_callback.py: await update_cache() instead of create_task 2. proxy_server.py: await async_set_cache_pipeline() instead of create_task 3. auth_checks.py: prefer valid_token.team_member_spend (from fresh cache) over team_membership.spend (which may be stale) This ensures budget checks see the most recent spend values and properly enforce budget limits when requests come in quick succession. Fixes: test_users_in_team_budget, test_chat_completion_low_budget --- litellm/proxy/auth/auth_checks.py | 8 +++++++- .../proxy/hooks/proxy_track_cost_callback.py | 20 +++++++++---------- litellm/proxy/proxy_server.py | 11 +++++----- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e0b056d450f..faca49ccc6e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2295,7 +2295,13 @@ async def _check_team_member_budget( and team_membership.litellm_budget_table.max_budget is not None ): team_member_budget = team_membership.litellm_budget_table.max_budget - team_member_spend = team_membership.spend or 0.0 + # Prefer valid_token.team_member_spend (from token cache) over team_membership.spend + # The token cache is updated synchronously after each request, while the team membership + # cache may have stale data since DB updates are batched + if valid_token.team_member_spend is not None: + team_member_spend = valid_token.team_member_spend + else: + team_member_spend = team_membership.spend or 0.0 if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index dab5fb1bfd5..8f681c727ea 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -181,17 +181,15 @@ class _ProxyDBLogger(CustomLogger): org_id=org_id, ) - # update cache - asyncio.create_task( - update_cache( - token=user_api_key, - user_id=user_id, - end_user_id=end_user_id, - response_cost=response_cost, - team_id=team_id, - parent_otel_span=parent_otel_span, - tags=tags, - ) + # update cache - await to ensure budget checks see updated spend + await update_cache( + token=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + response_cost=response_cost, + team_id=team_id, + parent_otel_span=parent_otel_span, + tags=tags, ) await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ffa2da4cfbe..91094c49c9f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1703,12 +1703,11 @@ async def update_cache( # noqa: PLR0915 if tags is not None: await _update_tag_cache() - asyncio.create_task( - user_api_key_cache.async_set_cache_pipeline( - cache_list=values_to_update_in_cache, - ttl=60, - litellm_parent_otel_span=parent_otel_span, - ) + # Await cache update to ensure budget checks see updated spend values + await user_api_key_cache.async_set_cache_pipeline( + cache_list=values_to_update_in_cache, + ttl=60, + litellm_parent_otel_span=parent_otel_span, )