diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1ef5b1072ed..6b5c4fcb5e7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/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 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d48c0aa51ee..0e4a65a030b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3d6f8ac32c1..2cef482c84c 100644 --- a/litellm/proxy/proxy_server.py +++ b/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) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index ba047c4e6e6..25e727e8bec 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/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 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 index ae04c6f3c70..79ff5bb3f60 100644 --- 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 @@ -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