fix(spend): stop feeding team-key spend into the user budget counters

When general_settings.apply_user_budget_to_team_keys is not True (the
default), the user budget check already skips team keys
(auth_checks.py) but the measurement side still counted their spend
against the owner: the Redis spend:user:{user_id} counter (post-call
increment and admission warm-check) and LiteLLM_UserTable.spend (DB
writer). A user carrying an active team use case saw their personal
key rejected 422 'User over budget' because the team's spend had
filled their personal envelope.

One shared predicate (is_team_key_exempt_from_user_budget, same
convention as budget_reservation.py) now drives all three sites:
- _iter_entity_counter_keys: spend:user: is not bound for team keys,
  so admission_counter_keys and post_call_counter_keys agree;
- _user_scope in increment_spend_counters: no user-counter increment;
- _update_user_db: the user's LiteLLM_UserTable row is not updated,
  so a reseed of the counter from that row (spend_counter_reseed.py)
  cannot reintroduce the team's spend. The global proxy-budget row
  and team/team_member/end_user counters are untouched.

With the toggle on, the historical fully-counted behavior is kept.

litellm/proxy/db/db_spend_update_writer.py

@@ -87,6 +87,7 @@ from litellm.proxy.spend_tracking.savings import (
     extract_cache_read_tokens,
     marks_gateway_injection,
 )
+from litellm.proxy.spend_tracking.spend_counter_batch import is_team_key_exempt_from_user_budget
 from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
 from litellm.repositories.prisma_protocols import BatchTable
 from litellm.responses.streaming_iterator import spawn_logging_task
@@ -1144,6 +1145,7 @@ class DBSpendUpdateWriter:
                 prisma_client=prisma_client,
                 litellm_proxy_budget_name=litellm_proxy_budget_name,
                 end_user_id=end_user_id,
+                team_id=team_id,
             ),
             payload=payload_copy,
         )
@@ -1269,14 +1271,18 @@ class DBSpendUpdateWriter:
         prisma_client: PrismaClient | None,
         litellm_proxy_budget_name: str | None,
         end_user_id: str | None = None,
+        team_id: str | None = None,
     ):
         """
-        - Update that user's row
+        - Update that user's row (skipped for team keys unless
+          ``apply_user_budget_to_team_keys`` is True: the Redis ``spend:user:`` counter applies
+          the same exemption, so a reseed of the counter from this row cannot reintroduce the
+          team's spend)
         - Update litellm-proxy-budget row (global proxy spend)
         """
         try:
             if prisma_client is not None:  # update
-                user_ids: Final = [user_id]
+                user_ids: Final = [user_id if not is_team_key_exempt_from_user_budget(team_id) else None]
                 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

@@ -919,6 +919,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
     PendingSpendIncrement,
     active_spend_counter_batch,
     forget_spend_counter,
+    is_team_key_exempt_from_user_budget,
     post_call_counter_keys,
     read_batched_spend_counter,
     record_spend_counter_value,
@@ -3724,6 +3725,12 @@ async def _increment_spend_counters_batched(
         )
 
     async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
+        # apply_user_budget_to_team_keys off: a team key's spend must not feed the user
+        # counter — the user budget check already skips team keys (auth_checks.py) and
+        # budget_reservation.py reserves nothing against spend:user: for them either, so
+        # counting here would fill the owner's personal envelope with the team's spend.
+        if is_team_key_exempt_from_user_budget(team_id):
+            return ()
         user_counter_key: Final = f"spend:user:{scope_user_id}"
         if user_counter_key in reserved_counter_keys:
             return ()

litellm/proxy/spend_tracking/spend_counter_batch.py

@@ -180,6 +180,20 @@ def release_spend_counter_batch() -> None:
         batch.close()
 
 
+def is_team_key_exempt_from_user_budget(team_id: object) -> bool:
+    """True when a team key's spend must stay off the user budget accounting.
+
+    The user budget check already skips team keys unless ``general_settings.
+    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
+    from litellm.proxy.proxy_server import general_settings
+
+    return general_settings.get("apply_user_budget_to_team_keys") is not True
+
+
 def _iter_entity_counter_keys(
     token: object,
     team_id: object,
@@ -189,14 +203,15 @@ def _iter_entity_counter_keys(
     end_user_id: object,
 ) -> Iterator[str]:
     """Only string ids name a counter; anything else (None, or an unresolved placeholder in synthetic
-    logging payloads) simply has no counter to bind."""
+    logging payloads) simply has no counter to bind. ``spend:user:`` is skipped for team keys unless
+    ``apply_user_budget_to_team_keys`` opts in, matching the user budget check."""
     if isinstance(token, str):
         yield f"spend:key:{token}"
     if isinstance(team_id, str):
         yield f"spend:team:{team_id}"
         if isinstance(user_id, str):
             yield f"spend:team_member:{user_id}:{team_id}"
-    if isinstance(user_id, str):
+    if isinstance(user_id, str) and not is_team_key_exempt_from_user_budget(team_id):
         yield f"spend:user:{user_id}"
     if isinstance(end_user_id, str):
         yield f"spend:end_user:{end_user_id}"

tests/unit/proxy/proxy_server/test_spend_counters.py

@@ -54,6 +54,14 @@ from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_tea
 from .conftest import normalize
 
 
+@pytest.fixture(autouse=True)
+def _apply_user_budget_to_team_keys(monkeypatch):
+    """These pins count a team key's spend against the user counter, which the proxy only does
+    when ``apply_user_budget_to_team_keys`` opts in; the default exemption is pinned in
+    tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py."""
+    monkeypatch.setitem(ps.general_settings, "apply_user_budget_to_team_keys", True)
+
+
 def _make_spend_counter_cache(
     *,
     redis_get_value=None,

tests/unit/proxy/spend_tracking/test_spend_counter_batch.py

@@ -36,6 +36,14 @@ TOKEN_KEYS = frozenset(
 )
 
 
+@pytest.fixture(autouse=True)
+def _apply_user_budget_to_team_keys(monkeypatch):
+    """These pins count a team key's spend against the user counter, which the proxy only does
+    when ``apply_user_budget_to_team_keys`` opts in; the default exemption is pinned in
+    test_team_key_user_spend_exemption.py."""
+    monkeypatch.setitem(ps.general_settings, "apply_user_budget_to_team_keys", True)
+
+
 class CountingRedis(RedisCache):
     def __init__(self, store: dict[str, object] | None = None, fail: bool = False) -> None:
         self.store: dict[str, object] = dict(store or {})

tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py

@@ -0,0 +1,175 @@
+"""The spend of a team key must not feed the user budget unless ``general_settings.
+apply_user_budget_to_team_keys`` is True: the user budget check already skips team keys
+(auth_checks.py), so measuring their spend against ``spend:user:`` and the user's
+``LiteLLM_UserTable.spend`` row would fill the owner's personal envelope with the team's
+spend (422 on their personal key). The same predicate drives the admission counter keys,
+the post-call Redis increment, and the DB writer, so a reseed of ``spend:user:`` from
+``LiteLLM_UserTable.spend`` cannot reintroduce team spend."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+import litellm
+import litellm.proxy.proxy_server as ps
+from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
+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,
+    post_call_counter_keys,
+)
+
+TEAM_TOKEN = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id=None)
+
+
+@pytest.fixture
+def toggle_off(monkeypatch):
+    """Default proxy setting: team keys are exempt from the user budget."""
+    monkeypatch.setitem(ps.general_settings, "apply_user_budget_to_team_keys", False)
+
+
+@pytest.fixture
+def toggle_on(monkeypatch):
+    monkeypatch.setitem(ps.general_settings, "apply_user_budget_to_team_keys", True)
+
+
+def _increment_spend_counters_setup(monkeypatch):
+    fake_cache = MagicMock()
+    fake_cache.in_memory_cache.get_cache = MagicMock(return_value=None)
+    fake_cache.redis_cache = MagicMock()
+    fake_cache.redis_cache.async_get_cache = AsyncMock(return_value=None)
+    recorded: dict[str, float] = {}
+
+    async def _record_pipeline(increment_list, **_):
+        results = []
+        for op in increment_list:
+            recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"]
+            results.append(recorded[op["key"]])
+        return results
+
+    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline)
+    fake_user_cache = MagicMock()
+    fake_user_cache.async_get_cache = AsyncMock(return_value=None)
+    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+    monkeypatch.setattr(ps, "prisma_client", None)
+    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))
+    return recorded
+
+
+async def _post_call_user_transaction(response_cost=5.0, user_id="u1", team_id="t1", end_user_id=None):
+    """Drive the DB-side construction the way ``update_database`` does and flush the queue."""
+    db_writer = DBSpendUpdateWriter()
+    db_writer.spend_update_queue = SpendUpdateQueue()
+    with patch.object(litellm, "max_budget", 0):
+        await db_writer._update_user_db(
+            response_cost=response_cost,
+            user_id=user_id,
+            prisma_client=MagicMock(),
+            litellm_proxy_budget_name=None,
+            end_user_id=end_user_id,
+            team_id=team_id,
+        )
+    return await db_writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
+
+
+# 1. Team key, toggle off: the user counter and the user DB row are left alone,
+#    the team counters still get their spend.
+@pytest.mark.asyncio
+async def test_team_key_toggle_off_never_touches_the_user_counter_or_user_row(toggle_off, monkeypatch):
+    recorded = _increment_spend_counters_setup(monkeypatch)
+
+    await ps.increment_spend_counters(token="hashed", team_id="t1", user_id="u1", response_cost=5.0)
+
+    assert recorded == {
+        "spend:key:hashed": 5.0,
+        "spend:team:t1": 5.0,
+        "spend:team_member:u1:t1": 5.0,
+    }
+
+    queued = await _post_call_user_transaction()
+    assert [u["entity_type"] for u in queued] == []  # no user_list transaction in the DB payload
+
+
+@pytest.mark.asyncio
+async def test_team_key_toggle_off_keeps_team_transactions_in_the_db_payload(toggle_off):
+    db_writer = DBSpendUpdateWriter()
+    db_writer.spend_update_queue = SpendUpdateQueue()
+    await db_writer._update_team_db(
+        response_cost=5.0, team_id="t1", user_id="u1", prisma_client=MagicMock()
+    )
+    queued = await db_writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
+    team_transactions = [u for u in queued if u["entity_type"] == Litellm_EntityType.TEAM]
+    assert [u["entity_id"] for u in team_transactions] == ["t1"]  # team transaction present
+    assert all(u["entity_type"] != Litellm_EntityType.USER for u in queued)
+
+
+# 2. Team key, toggle on: historical behavior, the user is fully counted.
+@pytest.mark.asyncio
+async def test_team_key_toggle_on_counts_the_user_everywhere(toggle_on, monkeypatch):
+    recorded = _increment_spend_counters_setup(monkeypatch)
+
+    await ps.increment_spend_counters(token="hashed", team_id="t1", user_id="u1", response_cost=5.0)
+
+    assert recorded["spend:user:u1"] == 5.0
+    assert recorded["spend:team:t1"] == 5.0
+
+    queued = await _post_call_user_transaction()
+    user_transactions = [u for u in queued if u["entity_type"] == Litellm_EntityType.USER]
+    assert [u["entity_id"] for u in user_transactions] == ["u1"]
+
+
+# 3. Personal key (no team_id), toggle off: unchanged.
+@pytest.mark.asyncio
+async def test_personal_key_toggle_off_is_unchanged(toggle_off, monkeypatch):
+    recorded = _increment_spend_counters_setup(monkeypatch)
+
+    await ps.increment_spend_counters(token="hashed", team_id=None, user_id="u1", response_cost=5.0)
+
+    assert recorded["spend:user:u1"] == 5.0
+
+    queued = await _post_call_user_transaction(team_id=None)
+    user_transactions = [u for u in queued if u["entity_type"] == Litellm_EntityType.USER]
+    assert [u["entity_id"] for u in user_transactions] == ["u1"]
+
+
+# 4. Admission: a team key's admission counter keys exclude spend:user: unless the toggle is on.
+@pytest.mark.asyncio
+async def test_admission_counter_keys_exclude_user_counter_for_team_keys(toggle_off):
+    without_toggle = admission_counter_keys(TEAM_TOKEN, end_user_id=None)
+    assert "spend:user:u1" not in without_toggle
+    assert {"spend:key:hashed", "spend:team:t1", "spend:team_member:u1:t1"} <= without_toggle
+
+    post_call_off = post_call_counter_keys(
+        token="hashed", team_id="t1", user_id="u1", org_id=None, end_user_id=None, tags=None, model_access_groups=None
+    )
+    assert "spend:user:u1" not in post_call_off
+
+
+@pytest.mark.asyncio
+async def test_admission_counter_keys_include_user_counter_when_toggle_is_on(toggle_on):
+    assert "spend:user:u1" in admission_counter_keys(TEAM_TOKEN, end_user_id=None)
+    assert "spend:user:u1" in post_call_counter_keys(
+        token="hashed", team_id="t1", user_id="u1", org_id=None, end_user_id=None, tags=None, model_access_groups=None
+    )
+
+
+# 5. Re-seed coherence: the user counter is reseeded from LiteLLM_UserTable.spend alone, and
+#    the DB writer (pinned above) never adds team spend to that row, so a reseed cannot
+#    reintroduce the team's spend into spend:user:.
+@pytest.mark.asyncio
+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
+    )
+
+    reseeded = await SpendCounterReseed.from_db(prisma_client, counter_key="spend:user:u1")
+
+    assert reseeded == 51.77
+    prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "u1"})