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