Merge 23293f4f0c into 40be7fa1a8

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

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)
 

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)

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}"

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,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