feat(budgets): enforce and accumulate per-team token budgets

Enforcement: team_max_tokens_check mirrors team_max_budget_check and runs
in the same common_checks gather. It reads usage from the cross-pod
tokens:team:{id} counter (Redis-first via get_current_tokens_used),
falling back to the tokens_used column loaded in auth, and raises
BudgetExceededError when usage exceeds max_tokens_budget. A team with no
token budget set is skipped, so existing teams are unchanged.

Accumulation: the cost callback bumps the tokens:team:{id} counter by the
request's total tokens (awaited in the same batch scope as the spend
counter, so the next request's auth check sees it), and the batched team
DB update increments tokens_used beside spend. The redis update buffer
carries team_token_list_transactions so a cross-pod flush keeps the
token delta.

litellm/proxy/auth/auth_checks.py

@@ -1404,6 +1404,11 @@ async def common_checks(
                     proxy_logging_obj=proxy_logging_obj,
                     valid_token=valid_token,
                 ),
+                team_max_tokens_check(
+                    team_object=team_object,
+                    proxy_logging_obj=proxy_logging_obj,
+                    valid_token=valid_token,
+                ),
                 _team_multi_budget_check(team_object=team_object),
                 _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None,
                 _team_soft_budget_check(
@@ -6518,6 +6523,66 @@ async def team_max_budget_check(
             )
 
 
+
+async def team_max_tokens_check(
+    team_object: LiteLLM_TeamTable | None,
+    valid_token: UserAPIKeyAuth | None,
+    proxy_logging_obj: ProxyLogging,
+) -> None:
+    """Reject a request when the team is over its cumulative token budget.
+
+    Mirrors ``team_max_budget_check``: dollars are a poor proxy for a token
+    allowance once a team mixes models at different prices, so a team can cap
+    tokens directly (``max_tokens_budget`` + ``budget_duration``). The usage is
+    read from the cross-pod ``tokens:team:{id}`` counter (Redis-first), falling
+    back to the ``tokens_used`` column loaded in auth. No budget set on the team
+    means the check is skipped, so existing teams are unaffected.
+    """
+    max_tokens_budget = getattr(team_object, "max_tokens_budget", None)
+    if team_object is None or max_tokens_budget is None:
+        return
+
+    from litellm.proxy.proxy_server import get_current_tokens_used
+
+    tokens_used: Final = await get_current_tokens_used(
+        counter_key=f"tokens:team:{team_object.team_id}",
+        fallback_tokens=team_object.tokens_used or 0,
+    )
+
+    if tokens_used > max_tokens_budget:
+        if valid_token:
+            call_info: Final = CallInfo(
+                token=valid_token.token,
+                spend=float(tokens_used),
+                max_budget=float(max_tokens_budget),
+                user_id=valid_token.user_id,
+                team_id=valid_token.team_id,
+                team_alias=valid_token.team_alias,
+                organization_id=valid_token.org_id,
+                event_group=Litellm_EntityType.TEAM,
+            )
+            asyncio.create_task(
+                proxy_logging_obj.budget_alerts(
+                    type="team_budget",
+                    user_info=call_info,
+                )
+            )
+
+        raise litellm.BudgetExceededError(
+            current_cost=float(tokens_used),
+            max_budget=float(max_tokens_budget),
+            message=(
+                f"Token budget has been exceeded! Team={team_object.team_id} "
+                f"Tokens used: {tokens_used}, Max tokens: {max_tokens_budget}"
+            ),
+            entity_type=Litellm_EntityType.TEAM.value,
+            entity_id=team_object.team_id,
+        )
+
+
+_team_max_tokens_check: Final = team_max_tokens_check
+
+
 _team_max_budget_check: Final = team_max_budget_check
 
 

litellm/proxy/db/db_spend_update_writer.py

@@ -1066,6 +1066,7 @@ class DBSpendUpdateWriter:
                 team_id=team_id,
                 user_id=user_id,
                 prisma_client=prisma_client,
+                tokens=int(payload.get("total_tokens") or 0),
             )
         except Exception:
             verbose_proxy_logger.debug(
@@ -1281,6 +1282,7 @@ class DBSpendUpdateWriter:
         team_id: str | None,
         user_id: str | None,
         prisma_client: PrismaClient | None,
+        tokens: int | None = None,
     ):
         try:
             if team_id is None or prisma_client is None:
@@ -1294,6 +1296,7 @@ class DBSpendUpdateWriter:
                     entity_type=Litellm_EntityType.TEAM,
                     entity_id=team_id,
                     response_cost=response_cost,
+                    tokens=tokens,
                 )
             )
 
@@ -2165,6 +2168,9 @@ class DBSpendUpdateWriter:
 
         ### UPDATE TEAM TABLE ###
         team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"]
+        team_token_list_transactions: Final = (
+            db_spend_update_transactions.get("team_token_list_transactions") or {}
+        )
         verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions)
         if team_list_transactions is not None and len(team_list_transactions.keys()) > 0:
             for i in range(n_retry_times + 1):
@@ -2174,12 +2180,21 @@ class DBSpendUpdateWriter:
                         async with transaction.batch_() as batcher:
                             # Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
                             for team_id, response_cost in sorted(team_list_transactions.items()):
+                                # Increment tokens_used beside spend for token budgets; 0 when
+                                # no token budget is in play, which leaves the column at its default.
+                                token_increment = team_token_list_transactions.get(team_id) or 0
                                 verbose_proxy_logger.debug(
-                                    "Updating spend for team id=%s by %s", team_id, response_cost
+                                    "Updating spend for team id=%s by %s (tokens +%s)",
+                                    team_id,
+                                    response_cost,
+                                    token_increment,
                                 )
+                                team_update_data: dict = {"spend": {"increment": response_cost}}
+                                if token_increment:
+                                    team_update_data["tokens_used"] = {"increment": token_increment}
                                 batcher.litellm_teamtable.update_many(  # 'update_many' prevents error from being raised if no row exists
                                     where={"team_id": team_id},
-                                    data={"spend": {"increment": response_cost}},
+                                    data=team_update_data,
                                 )
                     break
                 except Exception as e:

litellm/proxy/db/db_transaction_queue/redis_update_buffer.py

@@ -74,6 +74,7 @@ _SpendTransactionField: TypeAlias = Literal[
     "end_user_list_transactions",
     "key_list_transactions",
     "team_list_transactions",
+    "team_token_list_transactions",
     "team_member_list_transactions",
     "org_list_transactions",
     "org_member_list_transactions",
@@ -88,6 +89,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
     "end_user_list_transactions",
     "key_list_transactions",
     "team_list_transactions",
+    "team_token_list_transactions",
     "team_member_list_transactions",
     "org_list_transactions",
     "org_member_list_transactions",
@@ -957,6 +959,9 @@ class RedisUpdateBuffer:
             end_user_list_transactions=_merged_entity_transactions(list_of_transactions, "end_user_list_transactions"),
             key_list_transactions=_merged_entity_transactions(list_of_transactions, "key_list_transactions"),
             team_list_transactions=_merged_entity_transactions(list_of_transactions, "team_list_transactions"),
+            team_token_list_transactions=_merged_entity_transactions(
+                list_of_transactions, "team_token_list_transactions"
+            ),
             team_member_list_transactions=_merged_entity_transactions(
                 list_of_transactions, "team_member_list_transactions"
             ),

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -733,6 +733,29 @@ class _IncrementSpendCounters(Protocol):
     ) -> None: ...
 
 
+def _extract_total_tokens(kwargs: dict, completion_response: object) -> int:
+    """Total tokens for this request, for the team token-budget counter.
+
+    Prefers the StandardLoggingPayload (always populated for a billed call),
+    falling back to the response usage. Returns 0 when unknown, which leaves the
+    token counter untouched.
+    """
+    sl_object = kwargs.get("standard_logging_object") or {}
+    total = sl_object.get("total_tokens")
+    if total:
+        return int(total)
+    usage = getattr(completion_response, "usage", None)
+    if usage is None and isinstance(completion_response, dict):
+        usage = completion_response.get("usage")
+    if usage is not None:
+        total = getattr(usage, "total_tokens", None)
+        if total is None and isinstance(usage, dict):
+            total = usage.get("total_tokens")
+        if total:
+            return int(total)
+    return 0
+
+
 async def _update_database_and_spend_counters(
     proxy_logging_obj: "ProxyLogging",
     increment_spend_counters: _IncrementSpendCounters,
@@ -874,6 +897,27 @@ async def _update_database_and_spend_counters_in_batch(
             finally:
                 budget_reservation["finalized"] = True
         raise
+
+    # Mirror the spend counter for the team token budget: bump the cross-pod
+    # tokens:team:{id} counter so the next request's auth check sees this
+    # request's tokens before the batched DB flush of tokens_used lands.
+    # Awaited like the spend counter above; best-effort so a counter hiccup
+    # never fails a request that was already billed.
+    if team_id is not None:
+        total_tokens: Final = _extract_total_tokens(kwargs=kwargs, completion_response=completion_response)
+        if total_tokens:
+            from litellm.proxy.proxy_server import increment_spend_counter_cache
+
+            try:
+                await increment_spend_counter_cache(
+                    counter_key=f"tokens:team:{team_id}",
+                    increment=total_tokens,
+                )
+            except Exception:
+                verbose_proxy_logger.warning(
+                    "Failed to increment team token-budget counter for team_id=%s", team_id
+                )
+
     return True
 
 

litellm/proxy/proxy_server.py

@@ -3020,6 +3020,24 @@ async def get_current_spend(
     return current
 
 
+async def get_current_tokens_used(counter_key: str, fallback_tokens: int) -> int:
+    """Cross-pod cumulative token usage for a team's token budget.
+
+    Reads the Redis-backed ``tokens:team:{id}`` counter first (authoritative
+    across pods), the same way :func:`get_current_spend` reads the spend
+    counter, and falls back to the ``tokens_used`` column loaded in auth when
+    the counter is cold or Redis is unreachable. The counter is incremented
+    post-call in the cost callback and invalidated on the budget-reset tick, so
+    a read after a reset re-seeds from the freshly-zeroed column. Token budgets
+    tolerate the spend path's fallback-to-DB behavior for the same reason spend
+    does: the column is the authoritative value the reset job maintains.
+    """
+    value, _ = await read_spend_counter_cache_value(counter_key=counter_key)
+    if value is not None:
+        return int(value)
+    return int(fallback_tokens or 0)
+
+
 @with_service_target(SPEND_COUNTERS_TARGET)
 async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None:
     """Raise a counter that has fallen below the authoritative DB spend (e.g.