feat: multiple concurrent budget windows per API key and team (#24883)

[ishaan-berri]

• feat(proxy): add BudgetLimitEntry type and wire budget_limits into key/team models

• feat(schema): add budget_limits Json column to VerificationToken and TeamTable

• feat(migrations): add migration for budget_limits column on keys and teams

• feat(keys): initialize budget_limits windows with reset_at on key create/update

• feat(teams): initialize budget_limits windows with reset_at on team create/update

• feat(auth): add _virtual_key_multi_budget_check and _team_multi_budget_check

• feat(auth): call multi-budget checks from common_checks for keys and teams

• feat(proxy): increment per-window Redis spend counters after each request

• feat(budget): reset individual budget windows on schedule via reset_budget_job

• feat(ui): add hourly option to BudgetDurationDropdown

• feat(ui): add budget_limits field to KeyResponse type

• feat(ui): add Budget Windows editor to key edit view

• feat(ui): add Budget Windows editor to create key form

• fix(proxy): strip budget_limits=None before Prisma upsert to fix login 500

Prisma rejects nullable JSON fields (Json? without @default) when passed as Python None — it needs the field omitted entirely so the DB stores NULL via the column's nullable constraint. This was breaking /v2/login because the UI session key creation path hit the upsert with budget_limits=None.

• ui(key-edit): use antd InputNumber+Button for budget windows, add reset hints

• ui(create-key): use antd InputNumber+Button for budget windows, add reset hints

• docs(users): add multiple budget windows section with API + dashboard walkthrough

• fix: BudgetExceededError returns HTTP 429 instead of 400

• Add status_code=429 to BudgetExceededError class
• auth_exception_handler hardcoded code=400 → code=429

• fix: no-op else branch in multi-budget auth checks causes KeyError

• BudgetLimitEntry objects must be coerced via model_dump() not left as-is
• Move _virtual_key_multi_budget_check into common_checks (was asymmetric with _team_multi_budget_check which already lived there)

• fix: len() on JSON string returns char count not window count

Guard with isinstance check + json.loads() before iterating per-window Redis counters in increment_spend_counters

• fix: silent except:pass hides Redis reset failures in reset_budget_windows

Log Redis counter reset failures as warnings so they are observable

• test: add unit tests for multi-budget window enforcement

5 tests covering: no budget_limits passes, under budget passes, over hourly window raises 429, over monthly window raises 429, BudgetLimitEntry objects coerced without KeyError

• fix: key per-window counters stable across reorders (duration key, not index)

• fix: team+key per-window spend increments use duration key, not index

• fix: budget window reset uses duration key; log failures instead of swallowing

• refactor: extract BudgetWindowsEditor to shared component

• refactor: key_edit_view imports BudgetWindowsEditor from shared component

• refactor: create_key_button imports BudgetWindowsEditor from shared component

---

Relevant issues

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Type

🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test

Changes

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
0 out of 2 committers have signed the CLA.

❌ ishaan-berri
❌ yuneng-berri
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Apr 6, 2026 9:04pm

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing litellm_ishaan_april3 (f6911b2) with main (48d4dec)}

In general, unused imports should be removed from the import list to keep the codebase clean, avoid confusion, and possibly reduce bundle size. Since InputNumber is unused, the best fix is to delete InputNumber from the antd named import without altering any other imports or functionality.

Concretely, in ui/litellm-dashboard/src/components/organisms/create_key_button.tsx, locate the line:

import { Button as Button2, Form, Input, InputNumber, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";

and remove InputNumber from the destructuring list so that the remaining imports stay intact. No additional methods, imports, or definitions are required.

In general, the correct fix for unused imports is to remove the specific symbols from the import statement, keeping only the ones that are actually used. This improves readability and avoids misleading future maintainers.

In this file, the best fix is to edit the antd import on line 8 to drop Button as AntButton and InputNumber while preserving Form, Input, Select, Switch, and Tooltip as-is. No other lines or imports need to be changed, and no new code or behavior is required. The edit is localized to the single import line in ui/litellm-dashboard/src/components/templates/key_edit_view.tsx.

[greptile-apps]

Greptile Summary

This PR introduces multiple concurrent budget windows per API key and team — for example a $5/hour cap and a $100/month cap on the same key. The feature spans the full stack: a new budget_limits JSONB column (migration included), BudgetLimitEntry Pydantic type, per-window Redis counters incremented after each request, auth-time enforcement via new _virtual_key_multi_budget_check / _team_multi_budget_check checks, a scheduler job to reset expired windows, and a BudgetWindowsEditor UI component. It also fixes BudgetExceededError to return HTTP 429 instead of 500.

Key changes:

• litellm/proxy/auth/auth_checks.py — _virtual_key_multi_budget_check and _team_multi_budget_check enforce per-window limits at auth time
• litellm/proxy/proxy_server.py — increment_spend_counters atomically increments per-window Redis counters after each request
• litellm/proxy/common_utils/reset_budget_job.py — reset_budget_windows resets expired per-window counters on schedule; contains a timezone comparison bug (see inline comment)
• litellm/proxy/_types.py — BudgetLimitEntry type; budget_limits field added to GenerateRequestBase, TeamBase, and LiteLLM_VerificationToken
• litellm/exceptions.py + auth_exception_handler.py — BudgetExceededError now returns 429
• Migration adds budget_limits JSONB to both tables with IF NOT EXISTS (idempotent)
• Five focused unit tests cover the key auth-check path

Issue to address before merge:

• The _reset_expired_window helper compares a timezone-stripped local-time value against datetime.utcnow(), which produces incorrect reset timing when litellm_settings.timezone is configured to a non-UTC zone (windows reset too early for negative offsets, too late for positive offsets)."

Confidence Score: 4/5

PR is mostly safe to merge; one P1 timezone comparison bug in reset_budget_windows affects deployments using non-UTC litellm_settings.timezone.

All prior review concerns (N+1 query, inline imports, type inconsistency, React key, duplicate-duration) are noted but not newly introduced blocking issues. The one new P1 finding is the timezone stripping in _reset_expired_window that causes windows to reset at the wrong UTC time for non-UTC timezone configurations. This only affects users who have explicitly set litellm_settings.timezone to a non-UTC zone; UTC deployments (the majority) are unaffected. Score of 4 reflects this single actionable fix needed before merge.

litellm/proxy/common_utils/reset_budget_job.py (timezone comparison bug in _reset_expired_window)

Important Files Changed

Filename	Overview
litellm/proxy/auth/auth_checks.py	Adds _virtual_key_multi_budget_check and _team_multi_budget_check that iterate budget_limits windows, query per-window Redis counters via get_current_spend, and raise BudgetExceededError (429) if any window is exceeded; both wired into common_checks.
litellm/proxy/common_utils/reset_budget_job.py	Adds reset_budget_windows() and _reset_expired_window() to reset expired per-window counters on schedule; contains a timezone comparison bug where ISO strings with non-UTC offsets are stripped before comparison with datetime.utcnow().
litellm/proxy/proxy_server.py	Adds increment_spend_counters() atomically incrementing per-entity and per-window Redis spend counters after each request; introduces get_current_spend() Redis-first helper.
litellm/proxy/_types.py	Introduces BudgetLimitEntry Pydantic model and adds budget_limits to GenerateRequestBase, TeamBase (typed List[BudgetLimitEntry]) and LiteLLM_VerificationToken (typed List[dict]).
litellm/exceptions.py	Fixes BudgetExceededError to set self.status_code = 429, enabling the auth exception handler to return HTTP 429 for budget-exceeded errors.
litellm/proxy/auth/auth_exception_handler.py	Adds explicit isinstance(e, litellm.BudgetExceededError) branch raising ProxyException with code=429.
litellm/proxy/management_endpoints/key_management_endpoints.py	Stamps reset_at timestamps onto each budget_limits window and JSON-serialises them on key create/update before DB write.
litellm/proxy/management_endpoints/team_endpoints.py	Stamps reset_at onto each budget_limits window on team create/update via extracted _set_budget_reset_at() helper.
tests/test_litellm/proxy/auth/test_multi_budget_windows.py	Five async pytest unit tests for _virtual_key_multi_budget_check: covers empty limits, under-budget, over-daily-window (429 + message check), over-monthly-window, and BudgetLimitEntry coercion.
ui/litellm-dashboard/src/components/key_team_helpers/BudgetWindowsEditor.tsx	New shared React component for editing budget windows; four fixed duration options (1h/24h/7d/30d), InputNumber for dollar limit, remove button and reset hint per row.
litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql	Idempotent ALTER TABLE … ADD COLUMN IF NOT EXISTS … JSONB migration adding budget_limits to LiteLLM_VerificationToken and LiteLLM_TeamTable.

Sequence Diagram

sequenceDiagram
    participant Client
    participant AuthChecks
    participant Redis as Redis (spend_counter_cache)
    participant DB as PostgreSQL
    participant ResetJob as reset_budget_job

    Client->>AuthChecks: request with API key
    AuthChecks->>Redis: get_current_spend(spend:key:{token}:window:{duration})
    Redis-->>AuthChecks: window_spend
    alt window_spend >= max_budget
        AuthChecks-->>Client: 429 BudgetExceededError
    else all windows OK
        AuthChecks-->>Client: proceed
    end

    Client->>AuthChecks: (after response) increment_spend_counters
    AuthChecks->>Redis: async_increment_cache(spend:key:{token}:window:{duration}, cost)
    Redis-->>AuthChecks: ok

    ResetJob->>DB: find_many(budget_limits != null)
    DB-->>ResetJob: keys/teams with budget_limits
    loop for each expired window
        ResetJob->>Redis: async_set_cache(counter_key, 0.0)
        ResetJob->>DB: update(budget_limits with new reset_at)
    end

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'litellm_ishaan_april6' int..." | Re-trigger Greptile}

[greptile-apps]

Unbounded find_many + N+1 update() writes inside a loop

Both the key and team loops fetch the entire table with no take limit, and then issue one update() call per changed record inside the loop. Per the project DB rules, large result sets must be paginated and writes must be batched rather than N+1'd.

On an installation with thousands of keys that have budget_limits, this will:

1. Materialise the full key table into memory (OOM risk)
2. Fire one DB round-trip per changed window (deadlock / latency risk during the reset job)

Consider pagination via cursor-based take/cursor and batching the window update payloads, or writing a single UPDATE … SET … WHERE … via update_many with a JSON column expression, and apply the same fix to the teams block (lines 626–649).

[greptile-apps]

Inline import inside function body violates project style guide

The project's CLAUDE.md explicitly states: "Avoid imports within methods — place all imports at the top of the file (module-level). The only exception is avoiding circular imports where absolutely necessary."

This same pattern (from litellm.proxy.proxy_server import get_current_spend placed inside a function body) also appears in _team_multi_budget_check at line ~3237. Both should be moved to module-level, or — if a circular import genuinely exists — a # noqa: PLC0415 comment with a brief explanation should be added.

The same pattern also appears in reset_budget_job.py (from litellm.proxy.proxy_server import spend_counter_cache inside reset_budget_windows/_reset_budget_common), and from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time inside function bodies in key_management_endpoints.py and team_endpoints.py.

Context Used: CLAUDE.md (source)

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

Type inconsistency: LiteLLM_VerificationToken.budget_limits typed as Optional[List[dict]]

GenerateRequestBase and TeamBase declare budget_limits as Optional[List[BudgetLimitEntry]], giving Pydantic the schema to validate each entry. However, LiteLLM_VerificationToken (the read-path type populated from the DB cache) uses the weaker Optional[List[dict]], so no field-level validation is performed when a cached token is loaded and the budget_limits list is iterated in _virtual_key_multi_budget_check. Aligning both declarations to Optional[List[BudgetLimitEntry]] would make the type system the single source of truth and avoid the isinstance(window, dict) … window.model_dump() defensive branches scattered through the auth and spend paths.

[greptile-apps]

Array index used as React key

Using array position as the key prop means React reuses existing DOM nodes with shifted indices when a row is deleted from the middle of the list. This can produce stale field values or flickering after a remove. Consider assigning a stable id string (e.g. from a ref counter) to each window entry on creation and using that as the key instead.

General approach: Break the cycle by stopping reset_budget_job.py from importing proxy_server. Instead, obtain spend_counter_cache via dependency injection (passing it in) or via a neutral utility module that both sides can import without a cycle. Since we cannot modify other files, the safest change is to refactor reset_budget_windows to receive spend_counter_cache as an argument and remove the import from proxy_server.

Concrete best fix here:

• Change reset_budget_windows to accept spend_counter_cache as a parameter instead of importing it:

async def reset_budget_windows(self, spend_counter_cache) -> None:
    ...

• Remove the line from litellm.proxy.proxy_server import spend_counter_cache.
• Use the parameter spend_counter_cache everywhere the function currently uses the imported spend_counter_cache.

This keeps all existing behavior the same, assuming callers (e.g., proxy_server) already have access to spend_counter_cache and can pass it in. It also removes the direct dependency from reset_budget_job.py to proxy_server, breaking the cycle CodeQL reported.

Specific changes in litellm/proxy/common_utils/reset_budget_job.py:

• Around line 588: update the method signature to async def reset_budget_windows(self, spend_counter_cache) -> None:.
• Remove line 593 (the from litellm.proxy.proxy_server import spend_counter_cache import).
• No new imports or helpers are needed; spend_counter_cache is already used in the body and will be provided by the caller.

[greptile-apps]

Duplicate budget_duration silently double-counts per-window Redis spend

When two entries in budget_limits share the same budget_duration (e.g., two "24h" rows), async_increment_cache is called twice for the identical Redis counter key on every request. After a $3 request the counter holds $6, and the auth check in _virtual_key_multi_budget_check compares the inflated value against each entry's max_budget — effectively enforcing limits at half the actual spend.

The same flaw applies to the team block at lines 1822–1831.

There is no server-side validation preventing duplicates, and the UI's BudgetWindowsEditor defaults every new row to "24h" (so clicking "+ Add Budget Window" twice creates duplicates without any warning).

A Pydantic @model_validator on GenerateRequestBase (and TeamBase) that checks uniqueness of budget_duration within the list would catch this at the API boundary:

@model_validator(mode="after")
def check_unique_budget_durations(self) -> "GenerateRequestBase":
    if self.budget_limits:
        seen: set = set()
        for entry in self.budget_limits:
            duration = (
                entry.budget_duration
                if hasattr(entry, "budget_duration")
                else entry.get("budget_duration")
            )
            if duration in seen:
                raise ValueError(
                    f"budget_limits contains duplicate budget_duration '{duration}'. "
                    "Each duration must appear at most once."
                )
            seen.add(duration)
    return self

The same guard should also be added in prepare_key_update_data since budget_limits can arrive as raw dicts on the update path.

[greptile-apps]

Default duration "24h" on every new row enables accidental duplicate-duration creation

addWindow always appends { budget_duration: "24h", max_budget: null }. A user who clicks "+ Add Budget Window" twice without changing the first row's duration ends up with two identical budget_duration values. As described in the server-side comment, this causes the Redis counter for that window to be incremented twice per request, silently halving the effective budget limit.

Consider disabling the Select options that are already in use, or at minimum picking the first unused duration as the default:

const usedDurations = new Set(value.map((w) => w.budget_duration));
const availableOptions = BUDGET_WINDOW_OPTIONS.filter(
  (o) => !usedDurations.has(o.value)
);
const addWindow = () => {
  const nextDuration = availableOptions[0]?.value ?? "24h";
  onChange([...value, { budget_duration: nextDuration, max_budget: null }]);
};

And inside the Select for each row, mark already-used options as disabled to prevent switching to a duplicate duration.

[greptile-apps]

Timezone stripping causes incorrect reset timing for non-UTC configurations

When litellm_settings.timezone is set to a non-UTC zone, get_budget_reset_time returns a timezone-aware datetime whose ISO string encodes a UTC offset (e.g. "2026-04-07T00:00:00+05:30" for IST). Inside _reset_expired_window, the code strips that offset unconditionally:

reset_at = datetime.fromisoformat(
    reset_at_str.replace("Z", "+00:00")
).replace(tzinfo=None)   # drops offset, keeps LOCAL clock value
if reset_at > now:       # now = datetime.utcnow() (always UTC)
    return False

For a negative-offset zone (e.g. US/Eastern, UTC-5):

• Stored: "2026-04-07T00:00:00-05:00" → stripped naive: 2026-04-07T00:00:00
• datetime.utcnow() at 4:30 UTC = 2026-04-07T04:30:00
• Comparison: 00:00 > 04:30 → False → window resets 30 min early (correct reset is at 05:00 UTC)

For a positive-offset zone (e.g. Asia/Tokyo, UTC+9):

• Stored: "2026-04-07T00:00:00+09:00" → stripped naive: 2026-04-07T00:00:00
• datetime.utcnow() at 15:30 UTC = 2026-04-06T15:30:00
• Comparison: 2026-04-07T00:00:00 > 2026-04-06T15:30:00 → True → window is NOT reset even though its UTC equivalent (2026-04-06T15:00:00Z) is already in the past

The fix is to convert reset_at to UTC before the naive comparison:

from datetime import timezone as _timezone

reset_at = datetime.fromisoformat(
    reset_at_str.replace("Z", "+00:00")
)
# Normalise to naive-UTC before comparing with datetime.utcnow()
if reset_at.tzinfo is not None:
    reset_at = reset_at.astimezone(_timezone.utc).replace(tzinfo=None)
if reset_at > now:
    return False

Alternatively, change now to datetime.now(timezone.utc) and skip the .replace(tzinfo=None) call entirely, keeping the comparison timezone-aware throughout.