fix(spend): gate the post-cost user cache with the shared team-key exemption predicate

Follow-up to 55120c7483 ("fix(spend): stop feeding team-key spend into
the user budget counters") — audit r1 fixes:

- update_cache (proxy_server.py): the post-cost LiteLLM_UserTable cache
  writeback is now gated by is_team_key_exempt_from_user_budget. That
  cached object.spend is the fallback_spend of the user budget check
  (auth_checks.py) when Redis and the counters are cold, so a team key
  still inflated the owner's personal envelope there. MAJOR.
- auth_checks._user_max_budget_check reuses the shared predicate
  instead of duplicating the toggle inline, so check and measurement
  share one definition. Behavior unchanged for string team ids.
- is_team_key_exempt_from_user_budget is tightened to
  isinstance(team_id, str) — only string ids name a counter, a
  non-string placeholder never exempts (same philosophy as
  _iter_entity_counter_keys).
- _update_user_db builds its user_ids list conditionally instead of
  seeding a None sentinel.
- Test: the reseed test uses a neutral 50.0 instead of a real-case
  figure; new tests cover the user-cache gating (team key toggle
  off/on, personal key control) and the non-string team id.

litellm/proxy/auth/auth_checks.py

@@ -147,7 +147,10 @@ from litellm.proxy.guardrails.tool_name_extraction import (
 from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
 from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
-from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
+from litellm.proxy.spend_tracking.spend_counter_batch import (
+    SPEND_COUNTERS_TARGET,
+    is_team_key_exempt_from_user_budget,
+)
 from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
 from litellm.repositories.budget_repository import BudgetRepository
 from litellm.repositories.object_permission_repository import ObjectPermissionRepository
@@ -1379,8 +1382,7 @@ async def common_checks(
             # 4.1 personal budget
             if user_object is None or user_object.max_budget is None:
                 return
-            is_team_key: Final = team_object is not None and team_object.team_id is not None
-            if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True:
+            if team_object is not None and is_team_key_exempt_from_user_budget(team_object.team_id):
                 return
 
             from litellm.proxy.proxy_server import get_current_spend

litellm/proxy/db/db_spend_update_writer.py

@@ -1282,7 +1282,7 @@ class DBSpendUpdateWriter:
         """
         try:
             if prisma_client is not None:  # update
-                user_ids: Final = [user_id if not is_team_key_exempt_from_user_budget(team_id) else None]
+                user_ids: Final = [] if is_team_key_exempt_from_user_budget(team_id) else [user_id]
                 if litellm.max_budget > 0:  # track global proxy budget, if user set max budget
                     user_ids.append(litellm_proxy_budget_name)
 

litellm/proxy/proxy_server.py

@@ -4472,7 +4472,7 @@ async def update_cache(
     if response_cost is None:
         return
 
-    if user_id is not None:
+    if user_id is not None and not is_team_key_exempt_from_user_budget(team_id):
         _add_spend(user_id, LiteLLM_UserTable, response_cost)
     if end_user_id is not None:
         _add_spend(end_user_cache_key(end_user_id), LiteLLM_EndUserTable, response_cost)

litellm/proxy/spend_tracking/spend_counter_batch.py

@@ -187,8 +187,8 @@ def is_team_key_exempt_from_user_budget(team_id: object) -> bool:
     apply_user_budget_to_team_keys`` is True (auth_checks.py), so measuring a team key's
     spend against the owner's ``spend:user:`` counter would fill their personal envelope
     with their team's spend. Same convention as budget_reservation.py."""
-    if team_id is None:
-        return False
+    if not isinstance(team_id, str):
+        return False  # only string ids name a team counter; a placeholder never exempts
     from litellm.proxy.proxy_server import general_settings
 
     return general_settings.get("apply_user_budget_to_team_keys") is not True

tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py

@@ -15,12 +15,15 @@ import pytest
 
 import litellm
 import litellm.proxy.proxy_server as ps
-from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
 from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
 from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
 from litellm.proxy.spend_tracking.spend_counter_batch import (
     admission_counter_keys,
+    is_team_key_exempt_from_user_budget,
     post_call_counter_keys,
 )
 
@@ -166,10 +169,76 @@ async def test_admission_counter_keys_include_user_counter_when_toggle_is_on(tog
 async def test_user_counter_reseed_reads_only_the_user_row(toggle_off):
     prisma_client = MagicMock()
     prisma_client.db.litellm_usertable.find_unique = AsyncMock(
-        return_value=SimpleNamespace(spend=51.77)  # personal-key spend only, team spend never written here
+        return_value=SimpleNamespace(spend=50.0)  # personal-key spend only, team spend never written here
     )
 
     reseeded = await SpendCounterReseed.from_db(prisma_client, counter_key="spend:user:u1")
 
-    assert reseeded == 51.77
+    assert reseeded == 50.0
     prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "u1"})
+
+
+# 6. Post-cost user cache (update_cache): a team key must not inflate the cached
+#    LiteLLM_UserTable.spend — it is the fallback_spend of the user budget check
+#    (auth_checks.py) when Redis and the counters are cold.
+def _seeded_user_cache() -> UserApiKeyCache:
+    cache = UserApiKeyCache(in_memory_cache=InMemoryCache())
+    cache.set_cache(
+        key="u1",
+        value=LiteLLM_UserTable(user_id="u1", spend=10.0),
+        local_only=True,
+        model_type=LiteLLM_UserTable,
+    )
+    seeded = cache.get_cache(key="u1", local_only=True, model_type=LiteLLM_UserTable)
+    assert seeded is not None and seeded.spend == 10.0
+    return cache
+
+
+async def _update_cache_user(user_cache: UserApiKeyCache, team_id: str | None) -> None:
+    await ps.update_cache(
+        token=None,
+        user_id="u1",
+        end_user_id=None,
+        team_id=team_id,
+        response_cost=5.0,
+        parent_otel_span=None,
+        cache=user_cache,
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_cache_does_not_inflate_the_user_cache_for_team_keys(toggle_off):
+    cache = _seeded_user_cache()
+
+    await _update_cache_user(cache, team_id="t1")
+
+    user_obj = cache.get_cache(key="u1", local_only=True, model_type=LiteLLM_UserTable)
+    assert user_obj is not None
+    assert user_obj.spend == 10.0
+
+    # Control: the same call on a personal key (no team_id) still counts.
+    personal = _seeded_user_cache()
+    await _update_cache_user(personal, team_id=None)
+    personal_obj = personal.get_cache(key="u1", local_only=True, model_type=LiteLLM_UserTable)
+    assert personal_obj is not None
+    assert personal_obj.spend == 15.0
+
+
+@pytest.mark.asyncio
+async def test_update_cache_inflates_the_user_cache_when_toggle_is_on(toggle_on):
+    cache = _seeded_user_cache()
+
+    await _update_cache_user(cache, team_id="t1")
+
+    user_obj = cache.get_cache(key="u1", local_only=True, model_type=LiteLLM_UserTable)
+    assert user_obj is not None
+    assert user_obj.spend == 15.0
+
+
+# 7. Only string team ids can exempt the user budget: a non-string placeholder
+#    (None, or an unresolved value in synthetic payloads) never exempts — same
+#    philosophy as the counter keys, where only string ids name a counter.
+def test_non_string_team_id_never_exempts(toggle_off):
+    assert is_team_key_exempt_from_user_budget(123) is False
+    assert is_team_key_exempt_from_user_budget(None) is False
+    assert is_team_key_exempt_from_user_budget("t1") is True