feat(router): add separate ITPM/OTPM deployment rate limits by Sameerlite · Pull Request #31952 · BerriAI/litellm

Sameerlite · 2026-07-02T08:27:27Z

Summary

Add separate input/output tokens per minute (ITPM/OTPM) enforcement for router deployments when enforce_model_rate_limits is enabled
Reserve estimated input + max output tokens pre-call, reconcile actual usage post-call, and refund reservations on failure
Expose ITPM/OTPM limits and remaining usage via rate-limit headers and aggregate limits in model group info

Test plan

poetry run pytest tests/test_litellm/test_router/test_io_token_rate_limits.py -v
Verify proxy deployment with itpm/otpm on a model returns x-ratelimit-*-input-tokens / x-ratelimit-*-output-tokens headers
Confirm 429 when ITPM reservation exceeds configured limit

Note

Medium Risk
Changes core router rate-limit enforcement and proxy response handling; incorrect reservation/reconcile logic could cause false 429s or limit bypass, though extensive tests cover concurrency, refunds, and header behavior.

Overview
Adds separate ITPM/OTPM limits on router deployments (via itpm/otpm on model config), enforced when enforce_model_rate_limits is in optional_pre_call_checks. Pre-call logic atomically reserves estimated input tokens and max output tokens; success reconciles to actual usage (cached read tokens excluded from ITPM); failures refund reservations. Client-forged reservation metadata is stripped before enforcement

Extends ModelRateLimitingCheck to run IO limits in sync and async pre-call paths alongside legacy TPM/RPM, with refunds on 429 and reconcile/refund hooks on success/failure logging

Surfaces limits in responses: model group info aggregates ITPM/OTPM; get_remaining_model_group_usage emits x-ratelimit-*-input-tokens / *-output-tokens (merged with standard TPM/RPM headers when both apply). set_response_headers is refactored through shared helpers; dict-shaped and bare async-generator responses get _hidden_params via get_hidden_params_dict / HiddenParamsAsyncIteratorWrapper. The proxy strips _hidden_params from dict response bodies so internal routing metadata does not leak in JSON

IO-limited deployments still contribute TPM usage to routing counters for mixed model groups

^{Reviewed by Cursor Bugbot for commit 346abc7. Bugbot is set up for automated code reviews on this repo. Configure here.}

Support input/output tokens per minute on deployments via enforce_model_rate_limits, with reservation, reconciliation, refund on failure, and rate-limit headers. Co-authored-by: Cursor <cursoragent@cursor.com>

Drop unrelated Black reformatting from router.py and types/router.py so the PR only contains functional ITPM/OTPM changes. Co-authored-by: Cursor <cursoragent@cursor.com>

codecov · 2026-07-02T08:31:49Z

Codecov Report

❌ Patch coverage is 46.73913% with 294 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...utils/pre_call_checks/io_token_rate_limit_check.py	30.84%	222 Missing ⚠️
...er_utils/pre_call_checks/model_rate_limit_check.py	13.46%	45 Missing ⚠️
litellm/router.py	76.92%	21 Missing ⚠️
litellm/router_utils/add_retry_fallback_headers.py	93.15%	5 Missing ⚠️
litellm/proxy/common_request_processing.py	85.71%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

greptile-apps · 2026-07-02T08:32:23Z

Greptile Summary

This PR adds separate ITPM/OTPM (input/output tokens per minute) enforcement for router deployments, activated when enforce_model_rate_limits is enabled. Pre-call logic atomically reserves estimated input tokens against ITPM and max_tokens against OTPM; post-call reconciliation adjusts to actual usage; failures and retries cleanly refund reservations.

Core rate-limit logic (io_token_rate_limit_check.py): atomic increment-with-rollback for both ITPM and OTPM, minimal 1-token reservation fallback when tokenizer estimation fails, correct separation of billable input vs. output in reconciliation, and refund_stale_reservation_before_retry to clear stale reservations before a retry overwrites kwargs.
Router integration (router.py, model_rate_limit_check.py): IO reconciliation runs before the model_id guard so reservations are always returned; TPM tracking falls through for deployments that set both limit types; get_remaining_model_group_usage now emits both IO and TPM/RPM headers without an early return.
Response headers & proxy cleanup (add_retry_fallback_headers.py, common_request_processing.py): header-attachment helpers are unified into shared functions; dict-shaped responses gain _hidden_params support; _hidden_params is stripped from dict response bodies before they reach the client so internal routing metadata does not leak in JSON.

Confidence Score: 5/5

Safe to merge — all previously identified concerns have been addressed and the core reservation/reconcile/refund lifecycle is correctly implemented.

All issues raised in prior review rounds have been resolved: OTPM now uses the same atomic increment-with-rollback as ITPM; ITPM reconciliation tracks only billable input tokens; the model_id guard no longer blocks IO reconciliation; stale reservations are refunded synchronously before a retry overwrites kwargs; the minimal-slot fallback (1 instead of the full limit) prevents estimation failures from serializing concurrent requests; and both IO and TPM/RPM headers are emitted without an early-return race.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py	New module implementing ITPM/OTPM reservation, reconciliation, and refund logic. Atomic increment-with-rollback for both ITPM and OTPM, minimal-slot reservation when token estimation fails, correct separation of billable input vs output in reconciliation, and clean stash clearing on retries/failures.
litellm/router_utils/pre_call_checks/model_rate_limit_check.py	Extended to run IO reconciliation before the model_id guard so IO-only deployments always return reservations; falls through to TPM tracking when both limit types are set; failure path unconditionally refunds via async_io_token_refund_failure.
litellm/router.py	Adds refund_stale_reservation_before_retry call in _update_kwargs_with_deployment to synchronously reclaim any stashed reservation before a retry overwrites kwargs; get_remaining_model_group_usage now emits both IO and TPM/RPM headers without early return; set_response_headers refactored to shared helpers.
litellm/router_utils/add_retry_fallback_headers.py	Refactored into shared helpers; adds HiddenParamsAsyncIteratorWrapper for bare async generators; dict-shaped responses now supported for _hidden_params attachment.
litellm/proxy/common_request_processing.py	Strips _hidden_params from dict response bodies before returning to clients; adopts get_hidden_params_dict helper consistently; prevents internal routing metadata from leaking in JSON responses.
litellm/types/router.py	Adds itpm/otpm fields to GenericLiteLLMParams, LiteLLMParamsTypedDict, and ModelGroupInfo; adds ITPM/OTPM RouterCacheEnum values.
tests/test_litellm/test_router/test_io_token_rate_limits.py	Comprehensive new test file covering reservation+reconcile lifecycle, 429 enforcement, OTPM atomic concurrency, refund-on-failure, retry stash clearing, header generation, and mixed IO+TPM deployments.

_{Reviews (25): Last reviewed commit: "fix(io_token_rate_limit_check): use mode..." | Re-trigger Greptile}

veria-ai · 2026-07-02T08:38:34Z

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 15 · PR risk: 0/10

Address Greptile review on separate ITPM/OTPM deployment rate limits. - OTPM is now reserved atomically pre-call with rollback, matching the ITPM path, so concurrent requests can no longer overshoot the configured output limit before reconciliation - ITPM counts input tokens only; it no longer accumulates completion tokens, so the input-token limit and x-ratelimit-limit-input-tokens header describe input usage as their names imply - _read_reservation_from_kwargs only falls back to litellm_params.metadata when the top-level metadata channel is absent, so production requests carrying a litellm_params.metadata dict still reconcile and refund their reservation Adds regression tests for OTPM atomicity under concurrency, input-only ITPM enforcement, and reservation lookup when litellm_params.metadata is present.

Sameerlite · 2026-07-02T08:49:12Z

@greptileai

…eader The in-flight replay for x-ratelimit-remaining-input-tokens subtracted total tokens (input + output) instead of input tokens only, so clients saw remaining input quota understated by the completion token count on every response. Now consistent with the input-only ITPM counter.

Sameerlite · 2026-07-02T08:58:12Z

@greptileai

When a deployment configures itpm/otpm alongside tpm/rpm, the io-token path takes over and the tpm/rpm limits are not enforced. Log a warning the first time such a conflicting deployment is seen so the supersession is not silent, and document the mutual exclusivity. Post-call reconciliation now only trues up a counter that was actually reserved against, so the itpm/otpm keys are no longer incremented for deployments that never configured that limit.

Sameerlite · 2026-07-02T09:14:26Z

@greptileai

Post-call reconciliation now keys off the exact cache key stashed at pre-call time rather than one recomputed from the response-time minute. This fixes two issues: a request whose pre-call estimate was 0 now still writes its actual billable input to the ITPM counter (previously it was skipped, leaving the limit unenforceable for that request), and a call that finishes in a later minute reconciles against the minute it reserved against instead of pushing a negative delta into the next minute. Counters are only touched when their limit is configured.

Sameerlite · 2026-07-02T09:36:26Z

@greptileai

async_log_success_event gated IO reconciliation behind the model_id guard that only the TPM tracking path needs. Since reconciliation works entirely from the cache keys stashed in kwargs, a success event whose standard_logging_object lacks model_id would skip reconciliation and leave the reservation on the counter until the TTL expired, wasting quota. Route the IO path first.

Sameerlite · 2026-07-02T09:48:53Z

@greptileai

For ITPM/OTPM model groups the counter is incremented at reservation time (pre-call), so the remaining values returned by get_remaining_model_group_usage already account for the current request. Replaying the in-flight delta on top double-counted it and understated x-ratelimit-remaining-input/output-tokens by up to max_tokens on every response. Skip the delta for io-token groups; the legacy TPM/RPM replay path is unchanged.

Sameerlite · 2026-07-02T10:03:33Z

@greptileai

async_io_token_refund_failure and async_io_token_reconcile_success now clear the stashed reservation keys from the request metadata once done. Otherwise, on a model group mixing IO-limited and non-IO deployments, a failed IO call that retries on a non-IO fallback left the stale sentinel in the shared request metadata; the fallback's success handler would divert into IO reconciliation against the already-refunded key, driving the ITPM counter negative and skipping the non-IO deployment's TPM tracking.

Sameerlite · 2026-07-02T10:36:36Z

@greptileai

Missing usage was reconciled as zero and fully refunded the pre-call reservation, allowing limit bypass on repeated successful calls. Only adjust counters when usage is resolved from the response or standard logging fields; otherwise keep the reservation until TTL expires. Co-authored-by: Cursor <cursoragent@cursor.com>

Deployments with both itpm/otpm and tpm/rpm previously returned after the IO reservation and skipped RPM/TPM checks. Run both paths and refund the IO reservation only when RPM/TPM rejects after a successful reservation. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-07-03T12:09:27Z

bugbot run

Sameerlite · 2026-07-03T12:16:12Z

@greptileai

The early return after IO-token reconciliation in log_success_event and async_log_success_event skipped the TPM counter increment, so the tpm_key the pre-call check reads was never written and tpm_limit was never actually enforced on deployments that also configure itpm/otpm. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-07-03T12:28:54Z

@greptileai

usage/standard_logging_object entries carrying only total_tokens (no prompt/completion or input/output breakdown) were treated as resolved usage, resolving to (0, 0) and refunding the full reservation. Both _usage_is_present and the standard_logging_object fallback now require an actual input/output breakdown before reconciling, keeping the reservation otherwise. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-07-03T12:32:11Z

@greptileai

_reservation_value(0, limit) reserved the entire limit whenever token estimation failed (empty/unsupported input, tokenizer error), letting one such request claim the whole bucket and 429 every concurrent request to the deployment until it completed. Reserve 1 token instead so estimation failures no longer serialize traffic. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-07-03T12:44:31Z

@greptileai

On retry, set_io_token_rate_limit_request_kwargs clears reservation sentinels from the shared kwargs dict before a background failure handler can refund them, stranding the counter until TTL. Refund and clear any stale reservation in _update_kwargs_with_deployment before stripping sentinels for the next attempt. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-07-03T13:03:04Z

@greptileai

Sameerlite · 2026-07-03T13:22:21Z

@greptileai

… estimate; document sync-refund Redis ceiling Pass the deployment litellm_params.model to token_counter so it uses the model's native tokenizer instead of the generic fallback, narrowing the reservation over/under-estimate window between pre-call and post-call reconcile. Add a ponytail: comment to refund_stale_reservation_before_retry explaining the known ceiling: the synchronous DualCache.increment_cache issues a blocking Redis INCR when a Redis backend is configured. This only fires on streaming mid-stream retries (non-streaming failures await their failure handler before the retry picks a new deployment, leaving no sentinels to refund). Upgrade path: make _update_kwargs_with_deployment async. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-07-03T13:38:02Z

@greptileai fixed the minor conerns, please update the score to 5/5 in the main review

mateo-berri

Is the veria concern legit?

Sameerlite · 2026-07-03T20:34:12Z

It was a False positive statement

mateo-berri

LGTM; thanks

Sameerlite and others added 2 commits July 2, 2026 13:56

feat(router): add separate ITPM/OTPM deployment rate limits

e7033ad

Support input/output tokens per minute on deployments via enforce_model_rate_limits, with reservation, reconciliation, refund on failure, and rate-limit headers. Co-authored-by: Cursor <cursoragent@cursor.com>

chore(router): keep ITPM/OTPM diff minimal in router.py

5360c92

Drop unrelated Black reformatting from router.py and types/router.py so the PR only contains functional ITPM/OTPM changes. Co-authored-by: Cursor <cursoragent@cursor.com>

greptile-apps Bot reviewed Jul 2, 2026

View reviewed changes

Comment thread litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py Outdated

Comment thread litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py Outdated

Comment thread litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py Outdated