diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af60a315c60..087265d51aa 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -151,7 +151,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 @@ -1418,8 +1421,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 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ff57596232f..c5c1a92554f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -90,6 +90,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 @@ -1310,6 +1311,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, ) @@ -1432,14 +1434,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 = [] 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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c1af5d6f980..bbc2bc79e0b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -977,6 +977,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, @@ -3916,6 +3917,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 () @@ -4684,7 +4691,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) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7e0821e091e..25e727e8bec 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 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 + + 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..79ff5bb3f60 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_team_key_user_spend_exemption.py @@ -0,0 +1,244 @@ +"""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.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, +) + +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=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 == 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