From 55120c7483b35ba3fb9fe05e133ff520753e3ad6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Lavner?= Date: Thu, 29 Oct 2026 19:14:51 +0100 Subject: [PATCH] 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 | 10 +- litellm/proxy/proxy_server.py | 7 + .../spend_tracking/spend_counter_batch.py | 19 +- .../proxy/proxy_server/test_spend_counters.py | 8 + .../test_spend_counter_batch.py | 8 + .../test_team_key_user_spend_exemption.py | 175 ++++++++++++++++++ 6 files changed, 223 insertions(+), 4 deletions(-) create mode 100644 tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 20addef904b..d48c0aa51ee 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d4ec26614c0..3d6f8ac32c1 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 () diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7e0821e091e..ba047c4e6e6 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/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}" diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..9bbf99a6229 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/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, diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index 8ba79a45ff3..f4b7b1b79ef 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/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 {}) diff --git a/tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py b/tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py new file mode 100644 index 00000000000..ae04c6f3c70 --- /dev/null +++ b/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"})