Merge a1416d4f40 into c58fc0c996

litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql

@@ -0,0 +1,8 @@
+SET LOCAL lock_timeout = '10s';
+
+-- AlterTable
+-- Cumulative per-team token budget, mirroring the existing dollar budget
+-- (max_budget / spend). max_tokens_budget is the cap (null = untracked);
+-- tokens_used is the accumulator reset on the shared budget_reset_at tick.
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "max_tokens_budget" BIGINT;
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "tokens_used" BIGINT NOT NULL DEFAULT 0;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -197,6 +197,8 @@ model LiteLLM_TeamTable {
     tpd_limit     BigInt?
     budget_duration String? 
     budget_reset_at DateTime?
+    max_tokens_budget BigInt?  // cumulative token cap over budget_duration; null = untracked (mirrors max_budget)
+    tokens_used BigInt @default(0)  // cumulative tokens consumed in the current window (mirrors spend)
     blocked Boolean @default(false)
     created_at    DateTime               @default(now()) @map("created_at")
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")

litellm/models/team.py

@@ -74,6 +74,8 @@ class TeamBase(LiteLLMPydanticObjectBase):
     tpd_limit: int | None = None
     max_budget: float | None = None
     soft_budget: float | None = None
+    # cumulative token cap over budget_duration (mirrors max_budget); null = untracked
+    max_tokens_budget: int | None = None
     budget_duration: str | None = None
     budget_limits: list[BudgetLimitEntry] | None = None
     models: list[str] = []
@@ -86,6 +88,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
 class LiteLLM_TeamTable(TeamBase):
     team_id: str
     spend: float | None = None
+    tokens_used: int | None = None  # cumulative tokens consumed this window (mirrors spend)
     max_parallel_requests: int | None = None
     budget_duration: str | None = None
     budget_reset_at: datetime | None = None

litellm/proxy/_types.py

@@ -1958,6 +1958,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
     tpd_limit: int | None = None
     max_budget: float | None = None
     soft_budget: float | None = None
+    max_tokens_budget: int | None = None  # cumulative token cap over budget_duration
     models: list | None = None
     blocked: bool | None = None
     budget_duration: str | None = None
@@ -5179,6 +5180,9 @@ class DBSpendUpdateTransactions(TypedDict):
     end_user_list_transactions: dict[str, float] | None
     key_list_transactions: dict[str, float] | None
     team_list_transactions: dict[str, float] | None
+    # cumulative token deltas per team, flushed to LiteLLM_TeamTable.tokens_used
+    # alongside team_list_transactions (spend). Tokens budgets only apply to teams.
+    team_token_list_transactions: dict[str, float] | None  # mutable-ok: accumulator  # writable-ok: queue-reassigned
     team_member_list_transactions: dict[str, float] | None
     org_list_transactions: dict[str, float] | None
     org_member_list_transactions: ReadOnly[dict[str, float] | None]
@@ -5192,6 +5196,7 @@ class SpendUpdateQueueItem(TypedDict, total=False):
     entity_type: Litellm_EntityType
     entity_id: str
     response_cost: float | None
+    tokens: int | None  # writable-ok: total request tokens; dedup accumulates in place
 
 
 class ToolDiscoveryQueueItem(TypedDict, total=False):

litellm/proxy/auth/auth_checks.py

@@ -1407,6 +1407,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(
@@ -6908,6 +6913,67 @@ 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.
+    """
+    if team_object is None:
+        return
+    max_tokens_budget: Final = team_object.max_tokens_budget
+    if 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/common_utils/reset_budget_job.py

@@ -1400,6 +1400,9 @@ class ResetBudgetJob:
                         team_id = getattr(t.row, "team_id", None)
                         if team_id:
                             await self._invalidate_spend_counter(f"spend:team:{team_id}")
+                            # Drop the token-budget counter too, so the next read
+                            # re-seeds from the freshly-zeroed tokens_used column.
+                            await self._invalidate_spend_counter(f"tokens:team:{team_id}")
 
             end_time = time.time()
             outcome: Final = _ChunkOutcome(

litellm/proxy/db/db_spend_update_writer.py

@@ -1060,12 +1060,16 @@ class DBSpendUpdateWriter:
                 traceback.format_exc(),
             )
 
+        team_tokens: Final = payload.get(  # pyright: ignore[reportUnknownMemberType]  # SpendLogsPayload.get overload is partially unknown; total_tokens is a required int
+            "total_tokens", 0
+        )
         try:
             await self._update_team_db(
                 response_cost=response_cost,
                 team_id=team_id,
                 user_id=user_id,
                 prisma_client=prisma_client,
+                tokens=team_tokens,
             )
         except Exception:
             verbose_proxy_logger.debug(
@@ -1281,6 +1285,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 +1299,7 @@ class DBSpendUpdateWriter:
                     entity_type=Litellm_EntityType.TEAM,
                     entity_id=team_id,
                     response_cost=response_cost,
+                    tokens=tokens,
                 )
             )
 
@@ -2165,6 +2171,7 @@ 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 +2181,26 @@ 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 = (
+                                    {
+                                        "spend": {"increment": response_cost},
+                                        "tokens_used": {"increment": token_increment},
+                                    }
+                                    if token_increment
+                                    else {"spend": {"increment": response_cost}}
                                 )
                                 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/db/db_transaction_queue/spend_update_queue.py

@@ -121,6 +121,9 @@ class SpendUpdateQueue(BaseUpdateQueue):
                 current_cost = _in_memory_map[_key].get("response_cost", 0) or 0
                 update_cost = update.get("response_cost", 0) or 0
                 _in_memory_map[_key]["response_cost"] = current_cost + update_cost
+                current_tokens = _in_memory_map[_key].get("tokens", 0) or 0
+                update_tokens = update.get("tokens", 0) or 0
+                _in_memory_map[_key]["tokens"] = current_tokens + update_tokens
 
         for _key, update in _in_memory_map.items():
             aggregated_spend_updates.append(update)
@@ -138,6 +141,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
             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={},
@@ -219,6 +223,15 @@ class SpendUpdateQueue(BaseUpdateQueue):
 
             transactions_dict[entity_id] += response_cost or 0
 
+            # Tokens budgets only apply to teams, so sum the token delta into a
+            # parallel map flushed to LiteLLM_TeamTable.tokens_used beside spend.
+            if entity_type == Litellm_EntityType.TEAM:
+                team_tokens = db_spend_update_transactions["team_token_list_transactions"]
+                if team_tokens is None:
+                    team_tokens = {}
+                    db_spend_update_transactions["team_token_list_transactions"] = team_tokens
+                team_tokens[entity_id] = (team_tokens.get(entity_id) or 0) + (update.get("tokens") or 0)
+
         return db_spend_update_transactions
 
     async def _emit_new_item_added_to_queue_event(

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -382,6 +382,16 @@ class _ProxyDBLogger(CustomLogger):
                 router=get_llm_router(),
             )
 
+            # An unpriced model has no dollar cost and skips the spend pipeline below, but it
+            # still consumes tokens a team token budget must count, so advance the counter here.
+            if (
+                response_cost is None
+                and team_id is not None
+                and ProxyUpdateSpend.disable_spend_updates() is not True
+                and not _is_unbilled_interaction_response(completion_response)
+            ):
+                await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object))
+
             if response_cost is not None:
                 user_api_key: Final = metadata.get("user_api_key", None)
                 verbose_proxy_logger.debug(
@@ -735,6 +745,35 @@ class _IncrementSpendCounters(Protocol):
     ) -> None: ...
 
 
+def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int:
+    """Total tokens for this request, for the team token-budget counter.
+
+    Read from the StandardLoggingPayload, the same source the batched tokens_used
+    DB write uses, so the cross-pod counter and the column stay in agreement.
+    """
+    if sl_object is None:
+        return 0
+    return sl_object.get(  # pyright: ignore[reportUnknownMemberType]  # StandardLoggingPayload.get overload is partially unknown; total_tokens is a required int
+        "total_tokens", 0
+    )
+
+
+async def _increment_team_token_counter(team_id: str, total_tokens: int) -> None:
+    """Advance the cross-pod ``tokens:team:{id}`` counter so the next request's
+    auth check sees this request's tokens before the batched ``tokens_used``
+    flush lands. Best-effort, a counter hiccup must never fail a request that was
+    already served.
+    """
+    if not total_tokens:
+        return
+    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:  # noqa: BLE001  # best-effort cross-pod token counter; a Redis failure must not break the response or the spend write
+        verbose_proxy_logger.warning("Failed to increment team token-budget counter for team_id=%s", team_id)
+
+
 async def _update_database_and_spend_counters(
     proxy_logging_obj: "ProxyLogging",
     increment_spend_counters: _IncrementSpendCounters,
@@ -876,6 +915,19 @@ 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:
+        sl_object: Final = cast(  # cast-ok: typed shape of the untyped CustomLogger callback payload
+            "StandardLoggingPayload | None",
+            kwargs.get("standard_logging_object"),  # pyright: ignore[reportUnknownMemberType]  # untyped callback kwargs
+        )
+        await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object))
+
     return True
 
 

litellm/proxy/management_endpoints/scim/scim_transformations.py

@@ -165,7 +165,7 @@ class ScimTransformations:
             raise HTTPException(status_code=500, detail={"error": "No database connected"})
 
         if isinstance(team, dict):
-            team = LiteLLM_TeamTable(**team)
+            team = LiteLLM_TeamTable.model_validate(team)
 
         # Get team members with proper display names
         scim_members: Final[list[SCIMMember]] = []

litellm/proxy/management_endpoints/team_callback_endpoints.py

@@ -351,7 +351,7 @@ async def add_team_callbacks(
         # authenticated key holder could overwrite another team's logging
         # config (and read back the credentials they wrote).
         if not await get_team_access().allows(
-            user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
+            user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
         ):
             team_access_denied()
 
@@ -534,7 +534,7 @@ async def delete_team_callback(
         # deregister its callbacks, otherwise any authenticated key holder could
         # silence another team's observability integration.
         if not await get_team_access().allows(
-            user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
+            user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
         ):
             team_access_denied()
 
@@ -667,7 +667,7 @@ async def disable_team_logging(
         # team may disable its logging — otherwise any authenticated key
         # holder can silence audit logging for any team.
         if not await get_team_access().allows(
-            user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
+            user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
         ):
             team_access_denied()
 
@@ -808,7 +808,7 @@ async def get_team_callbacks(
         # (Langfuse / Langsmith / GCS). Only proxy admins / org admins /
         # team admins of THIS team may read them.
         if not await get_team_access().allows(
-            user_api_key_dict, LiteLLM_TeamTable(**_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
+            user_api_key_dict, LiteLLM_TeamTable.model_validate(_existing_team.model_dump()), TEAM_OR_ORG_ADMIN
         ):
             team_access_denied()
 

litellm/proxy/management_endpoints/team_endpoints.py

@@ -1481,6 +1481,7 @@ async def new_team(
     - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement.
     - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement.
     - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
+    - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
     - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
     - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
     - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
@@ -1801,9 +1802,8 @@ async def new_team(
             )
 
         ## ADD TO TEAM TABLE
-        complete_team_data: Final = LiteLLM_TeamTable(
-            **data_json,
-            model_id=_model_id,
+        complete_team_data: Final = LiteLLM_TeamTable.model_validate(
+            {**data_json, "model_id": _model_id},
         )
 
         # Set Management Endpoint Metadata Fields
@@ -2260,6 +2260,7 @@ async def update_team(
     - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
     - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit
     - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
+    - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
     - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
     - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
     - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.

litellm/proxy/proxy_server.py

@@ -3152,6 +3152,28 @@ 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 read is floored at
+    that column. An idle counter expires after its TTL and the next post-call
+    increment re-seeds it from only that request's delta, so a bare counter read
+    can sit below the recorded usage and let a team keep spending past the cap.
+    Flooring at the authoritative column closes that gap the same way
+    ``get_current_spend`` reconciles a stale-low counter against recorded spend,
+    and also covers a counter that missed an increment when its logging payload
+    failed to build.
+    """
+    value, _ = await read_spend_counter_cache_value(counter_key=counter_key)
+    floor: Final = int(fallback_tokens or 0)
+    if value is not None:
+        return max(int(value), floor)
+    return floor
+
+
 @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.

litellm/proxy/schema.prisma

@@ -197,6 +197,8 @@ model LiteLLM_TeamTable {
     tpd_limit     BigInt?
     budget_duration String? 
     budget_reset_at DateTime?
+    max_tokens_budget BigInt?  // cumulative token cap over budget_duration; null = untracked (mirrors max_budget)
+    tokens_used BigInt @default(0)  // cumulative tokens consumed in the current window (mirrors spend)
     blocked Boolean @default(false)
     created_at    DateTime               @default(now()) @map("created_at")
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")

litellm/repositories/unit_of_work.py

@@ -56,9 +56,15 @@ class TeamSpendResetWrites:
     table: BatchTable
 
     def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None:
+        data: Final = dict(_spend_reset_data(budget_reset_at, spend_decrement))
+        # Token budgets are team-only, so zero the cumulative tokens_used on the
+        # same tick that resets spend. A hard reset (vs spend's decrement) is
+        # acceptable at a budget boundary: a token allowance need not carry the
+        # handful of tokens that can race the reset the way billing-exact spend does.
+        data["tokens_used"] = 0
         self.table.update(
             where={"team_id": team_id},
-            data=_spend_reset_data(budget_reset_at, spend_decrement),
+            data=data,
         )
 
 

schema.prisma

@@ -197,6 +197,8 @@ model LiteLLM_TeamTable {
     tpd_limit     BigInt?
     budget_duration String? 
     budget_reset_at DateTime?
+    max_tokens_budget BigInt?  // cumulative token cap over budget_duration; null = untracked (mirrors max_budget)
+    tokens_used BigInt @default(0)  // cumulative tokens consumed in the current window (mirrors spend)
     blocked Boolean @default(false)
     created_at    DateTime               @default(now()) @map("created_at")
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")

tests/unit/proxy/auth/test_team_token_budget.py

@@ -0,0 +1,187 @@
+"""Tests for the per-team cumulative token budget check.
+
+team_max_tokens_check mirrors team_max_budget_check: a team can cap tokens
+directly (max_tokens_budget + budget_duration) because a dollar cap drifts once
+the team mixes models at different per-token prices. Usage is read from the
+cross-pod tokens:team:{id} counter (get_current_tokens_used), falling back to
+the tokens_used column loaded in auth.
+"""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+import litellm
+from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
+from litellm.models.team import LiteLLM_TeamTable
+from litellm.proxy.auth.auth_checks import team_max_tokens_check
+from litellm.proxy.utils import ProxyLogging
+
+
+def _proxy_logging_obj() -> ProxyLogging:
+    obj = ProxyLogging.__new__(ProxyLogging)
+
+    async def _noop(*args, **kwargs):
+        return None
+
+    obj.budget_alerts = _noop  # type: ignore[attr-defined]
+    return obj
+
+
+def _patch_tokens(value):
+    async def _effect(counter_key, fallback_tokens):
+        return value if value is not None else int(fallback_tokens or 0)
+
+    return patch(
+        "litellm.proxy.proxy_server.get_current_tokens_used",
+        new_callable=AsyncMock,
+        side_effect=_effect,
+    )
+
+
+@pytest.mark.asyncio
+async def test_no_budget_set_is_skipped():
+    """A team with no token budget is never blocked, so existing teams are unaffected."""
+    team = LiteLLM_TeamTable(team_id="t1", tokens_used=10**9)
+    token = UserAPIKeyAuth(token="hashed", team_id="t1")
+    with _patch_tokens(10**9) as get_tokens:
+        result = await team_max_tokens_check(
+            team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()
+        )
+    assert result is None
+    assert get_tokens.await_count == 0
+
+
+@pytest.mark.asyncio
+async def test_under_budget_passes():
+    team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=0)
+    token = UserAPIKeyAuth(token="hashed", team_id="t1")
+    with _patch_tokens(500_000) as get_tokens:
+        result = await team_max_tokens_check(
+            team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()
+        )
+    assert result is None
+    assert get_tokens.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_over_budget_raises():
+    team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=0)
+    token = UserAPIKeyAuth(token="hashed", team_id="t1")
+    with _patch_tokens(1_200_000):
+        with pytest.raises(litellm.BudgetExceededError) as exc:
+            await team_max_tokens_check(team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj())
+    assert exc.value.current_cost == 1_200_000
+    assert exc.value.max_budget == 1_000_000
+    assert exc.value.entity_type == Litellm_EntityType.TEAM.value
+
+
+@pytest.mark.asyncio
+async def test_cold_counter_falls_back_to_db_tokens_used():
+    """When the counter is cold (None), usage comes from the tokens_used column
+    loaded in auth; an over-budget column still blocks."""
+    team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=1_500_000)
+    token = UserAPIKeyAuth(token="hashed", team_id="t1")
+    with _patch_tokens(None):
+        with pytest.raises(litellm.BudgetExceededError):
+            await team_max_tokens_check(team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj())
+
+
+@pytest.mark.asyncio
+async def test_exactly_at_cap_passes():
+    """Enforcement is strictly greater-than, matching the dollar budget check."""
+    team = LiteLLM_TeamTable(team_id="t1", max_tokens_budget=1_000_000, tokens_used=0)
+    token = UserAPIKeyAuth(token="hashed", team_id="t1")
+    with _patch_tokens(1_000_000) as get_tokens:
+        result = await team_max_tokens_check(
+            team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()
+        )
+    assert result is None
+    assert get_tokens.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_get_current_tokens_used_prefers_the_cross_pod_counter():
+    """A populated tokens:team counter is authoritative over the DB column."""
+    from litellm.proxy.proxy_server import get_current_tokens_used
+
+    async def _counter(counter_key):
+        return 500_000, None
+
+    with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _counter):
+        used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=7)
+    assert used == 500_000
+
+
+@pytest.mark.asyncio
+async def test_get_current_tokens_used_falls_back_to_db_when_counter_cold():
+    """A cold counter (None) falls back to the tokens_used column loaded in auth."""
+    from litellm.proxy.proxy_server import get_current_tokens_used
+
+    async def _cold(counter_key):
+        return None, None
+
+    with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _cold):
+        used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=1_234)
+    assert used == 1_234
+
+
+def test_extract_total_tokens_reads_the_standard_logging_payload():
+    """The counter delta is the StandardLoggingPayload total, the DB write's source."""
+    from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens
+
+    assert _extract_total_tokens({"total_tokens": 4242}) == 4242
+
+
+def test_extract_total_tokens_is_zero_without_a_payload():
+    """No payload leaves the counter untouched rather than raising."""
+    from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens
+
+    assert _extract_total_tokens(None) == 0
+
+
+@pytest.mark.asyncio
+async def test_get_current_tokens_used_floors_at_recorded_db_usage():
+    """A counter that reseeded low after a TTL expiry, or missed an increment when its
+    logging payload failed, must not read below the tokens_used column loaded in auth,
+    otherwise a team slips past the cap once the counter sits under recorded usage."""
+    from litellm.proxy.proxy_server import get_current_tokens_used
+
+    async def _stale_low(counter_key):
+        return 200, None
+
+    with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _stale_low):
+        used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=1_100)
+    assert used == 1_100
+
+
+@pytest.mark.asyncio
+async def test_unpriced_model_advances_the_team_token_counter():  # test-quality-ok: no in-process counter state
+    """An unpriced model returns response_cost=None and skips the spend pipeline, but its
+    tokens must still count against the team, otherwise it bypasses max_tokens_budget. The
+    cross-pod counter is Redis-backed, so the issued increment is the only in-process signal."""
+    from datetime import datetime
+
+    from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
+
+    kwargs = {
+        "model": "custom-unpriced",
+        "call_type": "acompletion",
+        "litellm_params": {"metadata": {"user_api_key_team_id": "t1"}},
+        "standard_logging_object": {
+            "response_cost": None,
+            "response_cost_failure_debug_info": "missing custom price",
+            "request_tags": None,
+            "total_tokens": 321,
+        },
+        "stream": False,
+    }
+
+    with patch("litellm.proxy.proxy_server.increment_spend_counter_cache", new_callable=AsyncMock) as inc:
+        await ProxyDBLogger()._PROXY_track_cost_callback(
+            kwargs=kwargs,
+            completion_response=None,
+            start_time=datetime(2026, 1, 1),
+            end_time=datetime(2026, 1, 1),
+        )
+    inc.assert_awaited_once_with(counter_key="tokens:team:t1", increment=321)

tests/unit/proxy/common_utils/test_reset_budget_job.py

@@ -431,7 +431,9 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
     assert write["where"] == {"team_id": "tid-1"}
     assert write["data"]["spend"] == {"decrement": 500.0}
     assert write["data"]["budget_reset_at"] > now
-    assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
+    # Team resets also hard-zero the token-budget accumulator for the new window.
+    assert write["data"]["tokens_used"] == 0
+    assert set(write["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"}
 
 
 def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client):
@@ -663,7 +665,11 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
         assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}"
         assert writes[0]["where"] == where
         assert writes[0]["data"]["spend"] == {"decrement": decrement}
-        assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
+        # Teams carry the extra token-budget accumulator, zeroed for the new window.
+        expected_keys = {"spend", "budget_reset_at"}
+        if table_name == "team":
+            expected_keys = expected_keys | {"tokens_used"}
+        assert set(writes[0]["data"].keys()) == expected_keys
 
     # The budget tier's cascade rides the same batch machinery.
     assert _batch_writes(mock_prisma_client, "enduser") == [
@@ -4216,9 +4222,13 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c
     assert len(user_writes) == 1
     assert user_writes[0]["where"] == {"user_id": "user2"}
     assert len(team_writes) == 2
-    for c in key_writes + user_writes + team_writes:
+    for c in key_writes + user_writes:
         assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
         assert c["data"]["spend"] == {"decrement": pre_reset_spend[next(iter(c["where"].values()))]}
+    # Team writes additionally zero the token-budget accumulator for the new window.
+    for c in team_writes:
+        assert set(c["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"}
+        assert c["data"]["spend"] == {"decrement": pre_reset_spend[next(iter(c["where"].values()))]}
 
 
 @pytest.mark.asyncio
@@ -4274,7 +4284,7 @@ async def test_reset_budget_teams_partial_failure():
     team_writes = [c for c in batch_calls if c["table"] == "team"]
     assert len(team_writes) == 1
     assert team_writes[0]["where"] == {"team_id": "team2"}
-    assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
+    assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"}
     assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
 
     failure_hook_calls = (

tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py

@@ -360,3 +360,44 @@ async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue):
     )
     assert aggregated["user_list_transactions"]["user1"] == 200 * 0.5
     assert aggregated["key_list_transactions"]["key1"] == 300 * 1.0
+
+
+@pytest.mark.asyncio
+async def test_team_token_aggregation(spend_queue):
+    """Per-team token deltas sum into team_token_list_transactions, parallel to
+    the spend sum, so LiteLLM_TeamTable.tokens_used is flushed on the same batch
+    cadence as spend. Non-team updates never create token transactions."""
+    updates: list[SpendUpdateQueueItem] = [
+        {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team1", "response_cost": 0.5, "tokens": 100},
+        {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team1", "response_cost": 0.25, "tokens": 50},
+        {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team2", "response_cost": 1.0, "tokens": 400},
+        # a team update without a tokens field counts as 0 tokens, spend still applies
+        {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team2", "response_cost": 0.2},
+        # a user update carries no team tokens
+        {"entity_type": Litellm_EntityType.USER, "entity_id": "user1", "response_cost": 0.1, "tokens": 9},
+    ]
+    for update in updates:
+        await spend_queue.update_queue.put(update)
+
+    aggregated = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
+
+    # spend is unchanged by the token plumbing
+    assert aggregated["team_list_transactions"]["team1"] == 0.75
+    assert aggregated["team_list_transactions"]["team2"] == 1.2
+    # tokens sum per team
+    assert aggregated["team_token_list_transactions"]["team1"] == 150
+    assert aggregated["team_token_list_transactions"]["team2"] == 400
+    # no team token entry leaked from the user update
+    assert "user1" not in aggregated["team_token_list_transactions"]
+
+
+@pytest.mark.asyncio
+async def test_team_token_aggregation_absent_when_no_tokens(spend_queue):
+    """A plain team spend update (no tokens) leaves tokens_used untouched: the
+    token map has the team at 0, which the commit treats as no increment."""
+    await spend_queue.update_queue.put(
+        {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team1", "response_cost": 0.5}
+    )
+    aggregated = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
+    assert aggregated["team_list_transactions"]["team1"] == 0.5
+    assert aggregated["team_token_list_transactions"].get("team1", 0) == 0

tests/unit/proxy/memory/test_memory_endpoints.py

@@ -197,6 +197,7 @@ def _make_team(team_id: str, *, admin_user_ids: List[str]) -> Any:
         created_at=now,
         updated_at=now,
         spend=0.0,
+        tokens_used=0,
         model_spend="{}",
         model_max_budget="{}",
         admins=[],

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -17010,6 +17010,7 @@ export interface paths {
          *     - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement.
          *     - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement.
          *     - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
+         *     - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
          *     - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
          *     - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
          *     - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
@@ -17241,6 +17242,7 @@ export interface paths {
          *     - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
          *     - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit
          *     - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
+         *     - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
          *     - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
          *     - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
          *     - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
@@ -34132,6 +34134,8 @@ export interface components {
             max_budget?: number | null;
             /** Max Parallel Requests */
             max_parallel_requests?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /**
              * Members
              * @default []
@@ -34190,6 +34194,8 @@ export interface components {
             team_id: string;
             /** Team Member Permissions */
             team_member_permissions?: string[] | null;
+            /** Tokens Used */
+            tokens_used?: number | null;
             /** Tpd Limit */
             tpd_limit?: number | null;
             /** Tpm Limit */
@@ -35667,6 +35673,8 @@ export interface components {
             max_budget?: number | null;
             /** Max Parallel Requests */
             max_parallel_requests?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /**
              * Members
              * @default []
@@ -35725,6 +35733,8 @@ export interface components {
             team_id: string;
             /** Team Member Permissions */
             team_member_permissions?: string[] | null;
+            /** Tokens Used */
+            tokens_used?: number | null;
             /** Tpd Limit */
             tpd_limit?: number | null;
             /** Tpm Limit */
@@ -39016,6 +39026,8 @@ export interface components {
             guardrails?: string[] | null;
             /** Max Budget */
             max_budget?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /** Mcp Rpm Limit */
             mcp_rpm_limit?: {
                 [key: string]: number;
@@ -40306,6 +40318,8 @@ export interface components {
             guardrails?: string[] | null;
             /** Max Budget */
             max_budget?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /** Mcp Rpm Limit */
             mcp_rpm_limit?: {
                 [key: string]: number;
@@ -45973,6 +45987,8 @@ export interface components {
             max_budget?: number | null;
             /** Max Parallel Requests */
             max_parallel_requests?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /**
              * Members
              * @default []
@@ -46031,6 +46047,8 @@ export interface components {
             team_id: string;
             /** Team Member Permissions */
             team_member_permissions?: string[] | null;
+            /** Tokens Used */
+            tokens_used?: number | null;
             /** Tpd Limit */
             tpd_limit?: number | null;
             /** Tpm Limit */
@@ -46112,6 +46130,8 @@ export interface components {
             max_budget?: number | null;
             /** Max Parallel Requests */
             max_parallel_requests?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /**
              * Members
              * @default []
@@ -46175,6 +46195,8 @@ export interface components {
             team_id: string;
             /** Team Member Permissions */
             team_member_permissions?: string[] | null;
+            /** Tokens Used */
+            tokens_used?: number | null;
             /** Tpd Limit */
             tpd_limit?: number | null;
             /** Tpm Limit */
@@ -48452,6 +48474,8 @@ export interface components {
             guardrails?: string[] | null;
             /** Max Budget */
             max_budget?: number | null;
+            /** Max Tokens Budget */
+            max_tokens_budget?: number | null;
             /** Mcp Rpm Limit */
             mcp_rpm_limit?: {
                 [key: string]: number;