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] 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.