From f67717ed1d83f18308f649a080c5742b68ff2a5d Mon Sep 17 00:00:00 2001 From: octo-patch Date: Thu, 23 Apr 2026 09:58:12 +0800 Subject: [PATCH] fix: skip personal spend update for team key calls (fixes #26239) _update_user_db() was unconditionally incrementing LiteLLM_UserTable.spend for every call, including calls made with team keys. This caused user personal spend to accumulate all team key costs, leading to false BudgetExceededError on personal key calls after a Redis flush (when the budget check falls back to the polluted DB value). The fix adds a team_id guard to _update_user_db() that returns early when team_id is set, matching the pattern already used by _update_team_db(). The budget check in auth_checks.py was already correctly skipping personal budget enforcement for team key calls; this aligns the write path with it. --- litellm/proxy/db/db_spend_update_writer.py | 9 +++ .../proxy/db/test_db_spend_update_writer.py | 67 +++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 241b66bc0ae..4b57eb961b2 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -334,6 +334,7 @@ class DBSpendUpdateWriter: user_api_key_cache=user_api_key_cache, litellm_proxy_budget_name=litellm_proxy_budget_name, end_user_id=end_user_id, + team_id=team_id, ) except Exception: verbose_proxy_logger.debug( @@ -501,11 +502,19 @@ class DBSpendUpdateWriter: user_api_key_cache: DualCache, litellm_proxy_budget_name: Optional[str], end_user_id: Optional[str] = None, + team_id: Optional[str] = None, ): """ - Update that user's row - Update litellm-proxy-budget row (global proxy spend) """ + # Skip personal spend tracking for team key calls. + # Team spend is tracked separately via _update_team_db. + # Without this guard, LiteLLM_UserTable.spend accumulates team key + # costs and causes false BudgetExceededError on personal key calls + # when the Redis counter is absent (e.g. after a restart). + if team_id is not None: + return ## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id) if existing_user_obj is not None and isinstance(existing_user_obj, dict): diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b5f82ef04c5..bd423f6d966 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1428,3 +1428,70 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called() mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called() mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_user_db_skips_for_team_key_calls(): + """ + Test that _update_user_db does NOT enqueue a spend update when team_id is set. + + Regression test for: https://github.com/BerriAI/litellm/issues/26239 + Bug: team key calls were incrementing LiteLLM_UserTable.spend, causing + false BudgetExceededError on personal key calls after a Redis flush. + """ + db_writer = DBSpendUpdateWriter() + db_writer.spend_update_queue = AsyncMock() + + mock_prisma_client = MagicMock() + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + # Call with team_id set — should NOT enqueue any update + await db_writer._update_user_db( + response_cost=0.05, + user_id="user-123", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + litellm_proxy_budget_name="litellm-proxy-budget", + end_user_id=None, + team_id="team-abc", + ) + + db_writer.spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_user_db_runs_for_personal_key_calls(): + """ + Test that _update_user_db DOES enqueue a spend update when team_id is None. + + Companion to test_update_user_db_skips_for_team_key_calls — ensures the + guard doesn't accidentally suppress personal key spend tracking. + """ + db_writer = DBSpendUpdateWriter() + db_writer.spend_update_queue = AsyncMock() + + mock_prisma_client = MagicMock() + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + import litellm + + original_max_budget = litellm.max_budget + litellm.max_budget = 0 # disable global proxy budget tracking + + try: + # Call with team_id=None — should enqueue the user spend update + await db_writer._update_user_db( + response_cost=0.05, + user_id="user-123", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + litellm_proxy_budget_name="litellm-proxy-budget", + end_user_id=None, + team_id=None, + ) + finally: + litellm.max_budget = original_max_budget + + db_writer.spend_update_queue.add_update.assert_called_once()