From 5d06961b864febd2b3bed0337a23d3990e6e2e4e Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 15:37:44 -0500
Subject: [PATCH 01/17] feat(budgets): schema + queue plumbing for per-team
 token budgets

Mirror the existing dollar budget (max_budget/spend) with a cumulative
token budget: a nullable max_tokens_budget cap and a tokens_used
accumulator on LiteLLM_TeamTable, reusing budget_duration/budget_reset_at
for the window. A team with no max_tokens_budget set is untouched.

Thread an optional per-request token delta through the spend update
queue: SpendUpdateQueueItem.tokens, summed per team into a parallel
team_token_list_transactions map alongside the existing spend sum, so
tokens_used is flushed on the same batch cadence as spend rather than
one DB write per request.
---
 .../migration.sql                                   |  6 ++++++
 .../litellm_proxy_extras/schema.prisma              |  2 ++
 litellm/models/team.py                              |  3 +++
 litellm/proxy/_types.py                             |  5 +++++
 .../db/db_transaction_queue/spend_update_queue.py   | 13 +++++++++++++
 litellm/proxy/schema.prisma                         |  2 ++
 schema.prisma                                       |  2 ++
 7 files changed, 33 insertions(+)
 create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql

diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql
new file mode 100644
index 00000000000..9306edd7d75
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql
@@ -0,0 +1,6 @@
+-- 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;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index f3ebee52dcd..3bef9af3b9e 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/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")
diff --git a/litellm/models/team.py b/litellm/models/team.py
index 8edf10703b1..79b3a6c0617 100644
--- a/litellm/models/team.py
+++ b/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
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index b6c81f17761..17978256526 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1947,6 +1947,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
@@ -5111,6 +5112,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, int] | None
     team_member_list_transactions: dict[str, float] | None
     org_list_transactions: dict[str, float] | None
     org_member_list_transactions: ReadOnly[dict[str, float] | None]
@@ -5124,6 +5128,7 @@ class SpendUpdateQueueItem(TypedDict, total=False):
     entity_type: Litellm_EntityType
     entity_id: str
     response_cost: float | None
+    tokens: int | None  # total tokens for this request; summed into tokens_used for TEAM updates
 
 
 class ToolDiscoveryQueueItem(TypedDict, total=False):
diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py
index 2b8535cb113..2524996d42f 100644
--- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py
+++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py
@@ -118,6 +118,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)
@@ -135,6 +138,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={},
@@ -216,6 +220,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(
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index f3ebee52dcd..3bef9af3b9e 100644
--- a/litellm/proxy/schema.prisma
+++ b/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")
diff --git a/schema.prisma b/schema.prisma
index f3ebee52dcd..3bef9af3b9e 100644
--- a/schema.prisma
+++ b/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")

From 0f079789877515df55b7f205ba62aa5ee8b306eb Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 15:41:58 -0500
Subject: [PATCH 02/17] 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             | 65 +++++++++++++++++++
 litellm/proxy/db/db_spend_update_writer.py    | 19 +++++-
 .../redis_update_buffer.py                    |  5 ++
 .../proxy/hooks/proxy_track_cost_callback.py  | 44 +++++++++++++
 litellm/proxy/proxy_server.py                 | 18 +++++
 5 files changed, 149 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 56773789059..c0c1d12e9ee 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/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
 
 
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index da8a7fb0f15..5b9e878113c 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/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:
diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
index 898aee493df..60f28a18768 100644
--- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
+++ b/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"
             ),
diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 0938d3f35c3..d670855820c 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/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
 
 
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 089647c6297..24fba4a5635 100644
--- a/litellm/proxy/proxy_server.py
+++ b/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.

From a1dd80689686fd39910eeafbcda46837f6eaff03 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 15:43:52 -0500
Subject: [PATCH 03/17] feat(budgets): reset tokens_used on the team budget
 window tick

On the budget_reset_at tick that zeroes spend, hard-reset the team's
tokens_used and invalidate the tokens:team:{id} counter so the next read
re-seeds from the freshly-zeroed column. The reset is team-scoped (token
budgets only apply to teams), so key/user resets are untouched.
---
 litellm/proxy/common_utils/reset_budget_job.py | 3 +++
 litellm/repositories/unit_of_work.py           | 8 +++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py
index 61330d12ac3..4c5db3c8089 100644
--- a/litellm/proxy/common_utils/reset_budget_job.py
+++ b/litellm/proxy/common_utils/reset_budget_job.py
@@ -1336,6 +1336,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(
diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py
index 10f092c1cb2..37f2cfbeaef 100644
--- a/litellm/repositories/unit_of_work.py
+++ b/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 = 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,
         )
 
 

From da179eb2f24d92ac58df4f86b651fd9b7f3dfce2 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 15:45:20 -0500
Subject: [PATCH 04/17] test(budgets): cover per-team token budget aggregation
 and enforcement

- token aggregation: per-team token sums in team_token_list_transactions,
  spend sums unchanged, non-team updates create no token entries
- enforcement: untracked teams skipped, under-budget passes, over-budget
  raises BudgetExceededError, cold counter falls back to tokens_used,
  strictly greater-than at the cap (matching the dollar budget check)
---
 .../unit/proxy/auth/test_team_token_budget.py | 97 +++++++++++++++++++
 .../test_spend_update_queue.py                | 41 ++++++++
 2 files changed, 138 insertions(+)
 create mode 100644 tests/unit/proxy/auth/test_team_token_budget.py

diff --git a/tests/unit/proxy/auth/test_team_token_budget.py b/tests/unit/proxy/auth/test_team_token_budget.py
new file mode 100644
index 00000000000..b598771b32e
--- /dev/null
+++ b/tests/unit/proxy/auth/test_team_token_budget.py
@@ -0,0 +1,97 @@
+"""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 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 _mock_get_current_tokens_used(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",
+        _mock_get_current_tokens_used,
+    )
+
+
+@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):
+        await team_max_tokens_check(
+            team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()
+        )  # no raise
+
+
+@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):
+        await team_max_tokens_check(
+            team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()
+        )  # no raise
+
+
+@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):
+        await team_max_tokens_check(
+            team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj()
+        )  # no raise
diff --git a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py
index 43f1a820885..7a1bac7d156 100644
--- a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py
+++ b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py
@@ -291,3 +291,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

From a4816fb9f0fb842efe351701c571deb72a7fb732 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 15:46:46 -0500
Subject: [PATCH 05/17] chore(budgets): type team_token_list_transactions as
 float map

Match the sibling transaction maps and the _merged_entity_transactions
return type so strict type-checking passes; values stay integral at
runtime and the column is BIGINT.
---
 litellm/proxy/_types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 17978256526..483d9485ec5 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -5114,7 +5114,7 @@ class DBSpendUpdateTransactions(TypedDict):
     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, int] | None
+    team_token_list_transactions: dict[str, float] | None
     team_member_list_transactions: dict[str, float] | None
     org_list_transactions: dict[str, float] | None
     org_member_list_transactions: ReadOnly[dict[str, float] | None]

From cbf9c40b9b42b89a9d73bb209551fc0cce69e5c7 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 20:50:48 -0500
Subject: [PATCH 06/17] fix(proxy): green up CI for team token budgets

- ruff format auth_checks / db_spend_update_writer / proxy_track_cost_callback
- document max_tokens_budget on /team/new and /team/update (api-docs gate)
- regenerate dashboard schema.d.ts from the updated OpenAPI spec
- reset-budget tests: team resets now also zero tokens_used, so assert the
  extra key for team writes only (key/user/org payloads are unchanged)
- memory test fixture: pass tokens_used=0 when hand-building the Prisma team
  row (the column is NOT NULL @default(0), exactly like spend)
---
 litellm/proxy/auth/auth_checks.py             |  1 -
 litellm/proxy/db/db_spend_update_writer.py    |  4 +---
 .../proxy/hooks/proxy_track_cost_callback.py  |  4 +---
 .../management_endpoints/team_endpoints.py    |  2 ++
 .../common_utils/test_reset_budget_job.py     | 18 ++++++++++----
 .../proxy/memory/test_memory_endpoints.py     |  1 +
 ui/litellm-dashboard/src/lib/http/schema.d.ts | 24 +++++++++++++++++++
 7 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index c0c1d12e9ee..7cc6b685a2a 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -6523,7 +6523,6 @@ async def team_max_budget_check(
             )
 
 
-
 async def team_max_tokens_check(
     team_object: LiteLLM_TeamTable | None,
     valid_token: UserAPIKeyAuth | None,
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 5b9e878113c..7e5137b8ade 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -2168,9 +2168,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 {}
-        )
+        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):
diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index d670855820c..122437f6d80 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -914,9 +914,7 @@ async def _update_database_and_spend_counters_in_batch(
                     increment=total_tokens,
                 )
             except Exception:
-                verbose_proxy_logger.warning(
-                    "Failed to increment team token-budget counter for team_id=%s", team_id
-                )
+                verbose_proxy_logger.warning("Failed to increment team token-budget counter for team_id=%s", team_id)
 
     return True
 
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 544bc471ce5..ae6da2ea98e 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -1448,6 +1448,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.
@@ -2227,6 +2228,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.
diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py
index c17aed49014..26c86bfd38f 100644
--- a/tests/unit/proxy/common_utils/test_reset_budget_job.py
+++ b/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):
@@ -552,7 +554,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") == [
@@ -4108,9 +4114,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
@@ -4166,7 +4176,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 = (
diff --git a/tests/unit/proxy/memory/test_memory_endpoints.py b/tests/unit/proxy/memory/test_memory_endpoints.py
index dff0e80fa77..f0a0fa9a259 100644
--- a/tests/unit/proxy/memory/test_memory_endpoints.py
+++ b/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=[],
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index ada0bb2b8a2..38d762e5d31 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -16912,6 +16912,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.
@@ -17143,6 +17144,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.
@@ -33745,6 +33747,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 []
@@ -33803,6 +33807,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 */
@@ -35276,6 +35282,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 []
@@ -35334,6 +35342,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 */
@@ -38596,6 +38606,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;
@@ -39837,6 +39849,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;
@@ -45475,6 +45489,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 []
@@ -45533,6 +45549,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 */
@@ -45612,6 +45630,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 []
@@ -45675,6 +45695,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 */
@@ -47851,6 +47873,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;

From dc8b436ed53d5f9f99262413902729cd7d5d3a4a Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 21:00:21 -0500
Subject: [PATCH 07/17] fix(proxy): annotate intentional blind-except on token
 counter (BLE001)

The best-effort token-budget counter deliberately swallows any error so a
Redis failure can't break the response or the spend write. Mark it with
# noqa: BLE001 (matching the sibling budget-reservation handlers in this file)
so the ruff strict-rule ratchet stays at its merge-base ceiling.
---
 litellm/proxy/hooks/proxy_track_cost_callback.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 122437f6d80..1352719aaa6 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -913,7 +913,7 @@ async def _update_database_and_spend_counters_in_batch(
                     counter_key=f"tokens:team:{team_id}",
                     increment=total_tokens,
                 )
-            except Exception:
+            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)
 
     return True

From cb4561e27bb59e92c0a5ce6e63301ec89c440a2c Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:05:22 -0500
Subject: [PATCH 08/17] fix(proxy): satisfy type-discipline ratchet for
 token-budget code

Keep the custom LIT rule counts at their merge-base ceiling:
- _extract_total_tokens: single-assignment Final locals (no rebinding) and
  Mapping[str, object] for the read-only kwargs param
- team_update_data: build the payload functionally under a Final binding
- max_tokens_budget / reset data: declare Final
- DBSpendUpdateTransactions.team_token_list_transactions and
  SpendUpdateQueueItem.tokens: mark the in-place accumulators mutable-ok /
  writable-ok (they are reassigned by the queue, mirroring their spend siblings)
---
 litellm/proxy/_types.py                       |  4 +--
 litellm/proxy/auth/auth_checks.py             |  2 +-
 litellm/proxy/db/db_spend_update_writer.py    | 11 +++++---
 .../proxy/hooks/proxy_track_cost_callback.py  | 27 +++++++++----------
 litellm/repositories/unit_of_work.py          |  2 +-
 5 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 483d9485ec5..95c3244be6c 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -5114,7 +5114,7 @@ class DBSpendUpdateTransactions(TypedDict):
     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
+    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]
@@ -5128,7 +5128,7 @@ class SpendUpdateQueueItem(TypedDict, total=False):
     entity_type: Litellm_EntityType
     entity_id: str
     response_cost: float | None
-    tokens: int | None  # total tokens for this request; summed into tokens_used for TEAM updates
+    tokens: int | None  # writable-ok: total request tokens; dedup accumulates in place
 
 
 class ToolDiscoveryQueueItem(TypedDict, total=False):
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 7cc6b685a2a..cfeb13f3b3f 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -6537,7 +6537,7 @@ async def team_max_tokens_check(
     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)
+    max_tokens_budget: Final = getattr(team_object, "max_tokens_budget", None)
     if team_object is None or max_tokens_budget is None:
         return
 
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 7e5137b8ade..657456f37a9 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -2187,9 +2187,14 @@ class DBSpendUpdateWriter:
                                     response_cost,
                                     token_increment,
                                 )
-                                team_update_data: dict = {"spend": {"increment": response_cost}}
-                                if token_increment:
-                                    team_update_data["tokens_used"] = {"increment": token_increment}
+                                team_update_data: Final = (
+                                    {
+                                        "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=team_update_data,
diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 1352719aaa6..3c3a05353e3 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -733,27 +733,24 @@ class _IncrementSpendCounters(Protocol):
     ) -> None: ...
 
 
-def _extract_total_tokens(kwargs: dict, completion_response: object) -> int:
+def _extract_total_tokens(kwargs: Mapping[str, object], 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
+    sl_object: Final = kwargs.get("standard_logging_object") or {}
+    sl_total: Final = sl_object.get("total_tokens")
+    if sl_total:
+        return int(sl_total)
+    usage: Final = (
+        completion_response.get("usage")
+        if isinstance(completion_response, dict)
+        else getattr(completion_response, "usage", None)
+    )
+    usage_total: Final = usage.get("total_tokens") if isinstance(usage, dict) else getattr(usage, "total_tokens", None)
+    return int(usage_total) if usage_total else 0
 
 
 async def _update_database_and_spend_counters(
diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py
index 37f2cfbeaef..3f8e44022b2 100644
--- a/litellm/repositories/unit_of_work.py
+++ b/litellm/repositories/unit_of_work.py
@@ -56,7 +56,7 @@ class TeamSpendResetWrites:
     table: BatchTable
 
     def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None:
-        data = dict(_spend_reset_data(budget_reset_at, spend_decrement))
+        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

From f453c4250968d9a92f70499b214146b403777fe9 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:14:05 -0500
Subject: [PATCH 09/17] test(proxy): assert observable output in token-budget
 pass-path tests

The three allow-path tests only checked that no exception was raised, which
the test-quality ratchet (TQ001) rejects. Assert the check returns None and
that the cross-pod counter was consulted exactly when a budget is set (zero
reads when none is set, one when it is), so a flipped comparison or a dropped
skip now fails the test.
---
 .../unit/proxy/auth/test_team_token_budget.py | 39 ++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/tests/unit/proxy/auth/test_team_token_budget.py b/tests/unit/proxy/auth/test_team_token_budget.py
index b598771b32e..261958fa162 100644
--- a/tests/unit/proxy/auth/test_team_token_budget.py
+++ b/tests/unit/proxy/auth/test_team_token_budget.py
@@ -7,7 +7,7 @@ cross-pod tokens:team:{id} counter (get_current_tokens_used), falling back to
 the tokens_used column loaded in auth.
 """
 
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
@@ -29,12 +29,13 @@ def _proxy_logging_obj() -> ProxyLogging:
 
 
 def _patch_tokens(value):
-    async def _mock_get_current_tokens_used(counter_key, fallback_tokens):
+    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",
-        _mock_get_current_tokens_used,
+        new_callable=AsyncMock,
+        side_effect=_effect,
     )
 
 
@@ -43,20 +44,24 @@ 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):
-        await team_max_tokens_check(
+    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()
-        )  # no raise
+        )
+    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):
-        await team_max_tokens_check(
+    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()
-        )  # no raise
+        )
+    assert result is None
+    assert get_tokens.await_count == 1
 
 
 @pytest.mark.asyncio
@@ -65,9 +70,7 @@ async def test_over_budget_raises():
     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()
-            )
+            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
@@ -81,9 +84,7 @@ async def test_cold_counter_falls_back_to_db_tokens_used():
     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()
-            )
+            await team_max_tokens_check(team_object=team, valid_token=token, proxy_logging_obj=_proxy_logging_obj())
 
 
 @pytest.mark.asyncio
@@ -91,7 +92,9 @@ 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):
-        await team_max_tokens_check(
+    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()
-        )  # no raise
+        )
+    assert result is None
+    assert get_tokens.await_count == 1

From fbda977c5ee632d0fcb475511ad6e88b6f099d1f Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:31:07 -0500
Subject: [PATCH 10/17] fix(proxy): keep token-budget code within the
 basedpyright ceiling

Type the token lookups instead of reading untyped callback data:
- _extract_total_tokens takes a StandardLoggingPayload and returns its typed
  total_tokens, the same source the batched tokens_used DB write uses
- read max_tokens_budget off the typed LiteLLM_TeamTable, not getattr
- read total_tokens off the typed SpendLogsPayload, not .get
- drop Final on the loop-assigned team_update_data (illegal, and the gate
  exempts loop vars from the Final rule anyway)
The one remaining untyped boundary (the CustomLogger kwargs dict) is cast once
with a reason.
---
 litellm/proxy/auth/auth_checks.py             |  6 +++--
 litellm/proxy/db/db_spend_update_writer.py    |  4 +--
 .../proxy/hooks/proxy_track_cost_callback.py  | 27 ++++++++-----------
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index cfeb13f3b3f..54d7f7de155 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -6537,8 +6537,10 @@ async def team_max_tokens_check(
     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: Final = getattr(team_object, "max_tokens_budget", None)
-    if team_object is None or max_tokens_budget is None:
+    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
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 657456f37a9..98c788b6fc4 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -1066,7 +1066,7 @@ class DBSpendUpdateWriter:
                 team_id=team_id,
                 user_id=user_id,
                 prisma_client=prisma_client,
-                tokens=int(payload.get("total_tokens") or 0),
+                tokens=payload["total_tokens"],
             )
         except Exception:
             verbose_proxy_logger.debug(
@@ -2187,7 +2187,7 @@ class DBSpendUpdateWriter:
                                     response_cost,
                                     token_increment,
                                 )
-                                team_update_data: Final = (
+                                team_update_data = (
                                     {
                                         "spend": {"increment": response_cost},
                                         "tokens_used": {"increment": token_increment},
diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 3c3a05353e3..6144dc37e7a 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -733,24 +733,15 @@ class _IncrementSpendCounters(Protocol):
     ) -> None: ...
 
 
-def _extract_total_tokens(kwargs: Mapping[str, object], completion_response: object) -> int:
+def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> 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.
+    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.
     """
-    sl_object: Final = kwargs.get("standard_logging_object") or {}
-    sl_total: Final = sl_object.get("total_tokens")
-    if sl_total:
-        return int(sl_total)
-    usage: Final = (
-        completion_response.get("usage")
-        if isinstance(completion_response, dict)
-        else getattr(completion_response, "usage", None)
-    )
-    usage_total: Final = usage.get("total_tokens") if isinstance(usage, dict) else getattr(usage, "total_tokens", None)
-    return int(usage_total) if usage_total else 0
+    if sl_object is None:
+        return 0
+    return sl_object["total_tokens"]
 
 
 async def _update_database_and_spend_counters(
@@ -901,7 +892,11 @@ async def _update_database_and_spend_counters_in_batch(
     # 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)
+        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
+        )
+        total_tokens: Final = _extract_total_tokens(sl_object)
         if total_tokens:
             from litellm.proxy.proxy_server import increment_spend_counter_cache
 

From 4eeb70009abc5077495a1b2fa9535143588bdf22 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:40:10 -0500
Subject: [PATCH 11/17] fix(proxy): tolerate a SpendLogsPayload without
 total_tokens

Reading payload["total_tokens"] unconditionally aborted _batch_database_updates
with a KeyError when a (test or partial) payload omits the key, dropping the
rest of the batch. Guard the typed index with a membership check so a missing
count falls back to 0, staying tolerant like the old .get while keeping the
read type-clean.
---
 litellm/proxy/db/db_spend_update_writer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 98c788b6fc4..a1838c3d4d1 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -1066,7 +1066,7 @@ class DBSpendUpdateWriter:
                 team_id=team_id,
                 user_id=user_id,
                 prisma_client=prisma_client,
-                tokens=payload["total_tokens"],
+                tokens=payload["total_tokens"] if "total_tokens" in payload else 0,
             )
         except Exception:
             verbose_proxy_logger.debug(

From f949f1735b0654ded233a71b888f7c9bbddce9bd Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:48:36 -0500
Subject: [PATCH 12/17] fix(proxy): read total_tokens with .get default, not an
 in-check ternary

The membership-guarded index tripped SIM401 (prefer dict.get), but a plain
.get on the SpendLogsPayload TypedDict trips basedpyright's unknown-member
rule. Use .get with a 0 default (tolerant and SIM401-clean) and suppress just
the one overload-unknown diagnostic with a reason.
---
 litellm/proxy/db/db_spend_update_writer.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index a1838c3d4d1..a1769946680 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -1060,13 +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=payload["total_tokens"] if "total_tokens" in payload else 0,
+                tokens=team_tokens,
             )
         except Exception:
             verbose_proxy_logger.debug(

From fe6edbef4cd9257df4d68e0f4523c1d4f5363759 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:57:34 -0500
Subject: [PATCH 13/17] fix(proxy): tolerate a StandardLoggingPayload without
 total_tokens

Some call types (e.g. aretrieve_batch) build a payload without total_tokens,
so indexing it raised KeyError and tripped the cost-tracking error path. Read
it with a .get default of 0 like the DB-side write, suppressing just the one
overload-unknown diagnostic with a reason.
---
 litellm/proxy/hooks/proxy_track_cost_callback.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 6144dc37e7a..4ebc44630ed 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -741,7 +741,9 @@ def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int:
     """
     if sl_object is None:
         return 0
-    return sl_object["total_tokens"]
+    return sl_object.get(  # pyright: ignore[reportUnknownMemberType]  # StandardLoggingPayload.get overload is partially unknown; total_tokens is a required int
+        "total_tokens", 0
+    )
 
 
 async def _update_database_and_spend_counters(

From dc088d9813a4e2742a0a35dd072023dd452871ac Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Mon, 12 Oct 2026 09:10:03 -0500
Subject: [PATCH 14/17] refactor(proxy): build LiteLLM_TeamTable via
 model_validate at **dict sites

Adding max_tokens_budget and tokens_used to LiteLLM_TeamTable made the existing
LiteLLM_TeamTable(**dict) constructions each flag two more basedpyright
unknown-argument / argument-type errors, which the type-check ratchet forbids.
model_validate takes the untyped dict as a single validated argument, so no
per-field diagnostics fire. It is also the correct pydantic v2 way to build a
model from a dict, and it leaves these call sites below their merge-base count.
---
 .../management_endpoints/scim/scim_transformations.py     | 2 +-
 .../proxy/management_endpoints/team_callback_endpoints.py | 8 ++++----
 litellm/proxy/management_endpoints/team_endpoints.py      | 5 ++---
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py
index 6a67093fde8..26d48d96845 100644
--- a/litellm/proxy/management_endpoints/scim/scim_transformations.py
+++ b/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]] = []
diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py
index 3908ab3c927..9d05c431c45 100644
--- a/litellm/proxy/management_endpoints/team_callback_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py
@@ -341,7 +341,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()
 
@@ -516,7 +516,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()
 
@@ -649,7 +649,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()
 
@@ -790,7 +790,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()
 
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index ae6da2ea98e..db3c31bd39d 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -1769,9 +1769,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

From 78a7fd1288b9075e5abf04dc60882f39eded0eb3 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Mon, 12 Oct 2026 10:21:41 -0500
Subject: [PATCH 15/17] test(proxy): cover get_current_tokens_used and
 _extract_total_tokens

Exercise the counter-read helper directly (Redis-first hit, cold-counter DB
fallback) and the token extractor (payload total, and 0 when the payload is
absent). The enforcement tests mock get_current_tokens_used, so these hit the
real code paths and lift patch coverage.
---
 .../unit/proxy/auth/test_team_token_budget.py | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/tests/unit/proxy/auth/test_team_token_budget.py b/tests/unit/proxy/auth/test_team_token_budget.py
index 261958fa162..361b42adebd 100644
--- a/tests/unit/proxy/auth/test_team_token_budget.py
+++ b/tests/unit/proxy/auth/test_team_token_budget.py
@@ -98,3 +98,43 @@ async def test_exactly_at_cap_passes():
         )
     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

From 7aafa3e4bd888e1f3502d3e7256958b0e8202db5 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Thu, 15 Oct 2026 09:34:11 -0500
Subject: [PATCH 16/17] fix(proxy): harden team token budget enforcement

Floor get_current_tokens_used at the recorded tokens_used column so a counter
that reseeded low after a TTL expiry, or missed an increment when its logging
payload failed, cannot read below real usage and let a team slip past the cap.
This mirrors how get_current_spend reconciles a stale-low counter against
recorded spend

Also advance the team token counter for successful unpriced responses
(response_cost None) that skip the spend pipeline, so an unpriced model no
longer bypasses max_tokens_budget
---
 .../proxy/hooks/proxy_track_cost_callback.py  | 38 ++++++++++-----
 litellm/proxy/proxy_server.py                 | 18 ++++---
 .../unit/proxy/auth/test_team_token_budget.py | 47 +++++++++++++++++++
 3 files changed, 85 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 4ebc44630ed..33fe9e4d54b 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -380,6 +380,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(
@@ -746,6 +756,22 @@ def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int:
     )
 
 
+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,
@@ -898,17 +924,7 @@ async def _update_database_and_spend_counters_in_batch(
             "StandardLoggingPayload | None",
             kwargs.get("standard_logging_object"),  # pyright: ignore[reportUnknownMemberType]  # untyped callback kwargs
         )
-        total_tokens: Final = _extract_total_tokens(sl_object)
-        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:  # 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)
+        await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object))
 
     return True
 
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 24fba4a5635..84ddc3d77b5 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -3026,16 +3026,20 @@ async def get_current_tokens_used(counter_key: str, fallback_tokens: int) -> int
     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.
+    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 int(value)
-    return int(fallback_tokens or 0)
+        return max(int(value), floor)
+    return floor
 
 
 @with_service_target(SPEND_COUNTERS_TARGET)
diff --git a/tests/unit/proxy/auth/test_team_token_budget.py b/tests/unit/proxy/auth/test_team_token_budget.py
index 361b42adebd..30080a31ee2 100644
--- a/tests/unit/proxy/auth/test_team_token_budget.py
+++ b/tests/unit/proxy/auth/test_team_token_budget.py
@@ -138,3 +138,50 @@ def test_extract_total_tokens_is_zero_without_a_payload():
     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)

From a1416d4f40ab9f50345d9e6432d0460fa1cb1482 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Thu, 15 Oct 2026 11:37:11 -0500
Subject: [PATCH 17/17] fix(migrations): set lock_timeout on team token budget
 migration

Upstream main added check_migrations_lock_timeout, which requires every
new migration's first statement to be SET LOCAL lock_timeout = '10s' so a
blocked ALTER cannot stall proxy boot indefinitely.
---
 .../20261011000000_add_team_token_budget/migration.sql          | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql
index 9306edd7d75..852414ff4b7 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261011000000_add_team_token_budget/migration.sql
@@ -1,3 +1,5 @@
+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);