feat(proxy): add project-level model_itpm_limit/model_otpm_limit - #35058
feat(proxy): add project-level model_itpm_limit/model_otpm_limit#35058shivijain2323 wants to merge 13 commits into
Conversation
We already support model_tpm_limit/model_rpm_limit per project, but that treats input and output tokens as one combined bucket, which doesn't match how Bedrock Mantle actually meters usage. This adds model_itpm_limit/model_otpm_limit so a project can cap input and output tokens separately, same shape as the existing fields. There's already deployment-level itpm/otpm support in router_utils/pre_call_checks/io_token_rate_limit_check.py, but that code only ever deals with one deployment at a time, whereas the proxy hook (parallel_request_limiter_v3.py) has to juggle key/team/org/project limits together on the same request. Rather than bolt a second copy of the reservation logic on, this just adds project-scoped ITPM/OTPM as two more descriptors into the same reservation path the hook already uses for TPM/RPM. If a project sets both the combined TPM limit and the new itpm/otpm limits on the same model, both get enforced, not one overriding the other, matching what the deployment-level check already does. Input and output are reserved separately up front, and if one side reserves fine but the other is over limit, we roll back the side that already went through so nothing gets left over-counted. On the way back, cached prompt tokens get excluded from the ITPM count since that's how Bedrock bills it, but that's purely for the rate limit math, cost and usage logging still see the full token count.
…ection async_post_call_failure_hook had its own refund path, separate from async_log_failure_event, for requests rejected before the LLM call ever ran (e.g. a downstream guardrail block). That path predates itpm/otpm and only knew about the combined TPM reservation: it refunded every token descriptor by the flat combined amount, which is wrong once the project itpm/otpm descriptors are in that same list, and it returned early whenever the combined reservation was zero, which is exactly the case when a project only sets model_itpm_limit/model_otpm_limit with no model_tpm_limit at all.
… io token reservation
… token reservation
async_release_max_parallel_requests_on_disconnect only released the parallel slot; any stashed ITPM/OTPM reservations stayed inflated until the window TTL, letting a caller repeatedly start and cancel streams to exhaust the project quota without consuming tokens. The fix mirrors the refund path already in async_log_failure_event: check the shared release guard, build pipeline ops to zero out each bucket, flush, then mark released so failure callbacks can't double-refund if they somehow fire later. The slot release is now conditional (if acquisition is not None) rather than an early-return gate, so the refund runs even on requests where the parallel slot was already cleared. Co-authored-by: Cursor <cursoragent@cursor.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR adds project-scoped input and output token-per-minute limits.
Confidence Score: 3/5This PR is not yet safe to merge because zero-output Responses requests can reserve no project OTPM while the dispatched request permits up to 16 output tokens. The limiter preserves max_output_tokens=0 and skips the reservation floor, but the downstream OpenAI Responses transformation normalizes that value to 16, allowing concurrent requests to generate unreserved output before post-call reconciliation. Files Needing Attention: litellm/proxy/hooks/parallel_request_limiter_v3.py, litellm/llms/openai/responses/transformation.py, tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
|
| Filename | Overview |
|---|---|
| litellm/proxy/hooks/parallel_request_limiter_v3.py | Adds project ITPM/OTPM reservation and reconciliation, but explicit zero-output Responses requests remain under-reserved because downstream normalization permits 16 output tokens. |
| litellm/models/project.py | Exposes the new dedicated project rate-limit columns through merged project metadata. |
| litellm/proxy/auth/user_api_key_auth.py | Propagates merged project metadata through the existing authentication paths. |
| litellm/proxy/_types.py | Adds the project ITPM/OTPM fields to project create and update request models. |
| litellm-proxy-extras/litellm_proxy_extras/migrations/20260722200000_add_project_itpm_otpm_limit/migration.sql | Adds non-null JSONB columns for project model ITPM and OTPM limits with empty-object defaults. |
| tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py | Expands limiter tests for split input/output reservations, refunds, reconciliation, and Responses API behavior. |
| tests/test_litellm/proxy/hooks/test_tpm_concurrent.py | Adds concurrency and zero-cap coverage, but the zero-cap expectation does not account for downstream normalization to 16 output tokens. |
Reviews (10): Last reviewed commit: "fix(proxy): honor an explicit zero outpu..." | Re-trigger Greptile
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| self._build_reservation_aware_tpm_ops( | ||
| targets=list(itpm_scopes), | ||
| reserved_scopes=itpm_scopes, | ||
| actual_tokens=0, |
There was a problem hiding this comment.
Medium: Stream cancellation can refund consumed tokens
Disconnect cleanup assumes actual usage is zero and refunds the full ITPM/OTPM reservation. The caller reaches this path whenever disconnect-time success logging is unavailable—for example when disable_streaming_logging is enabled or partial response assembly fails—even if chunks were already delivered. A project key can repeatedly read most of a stream and disconnect to consume tokens without retaining any quota charge. Keep the reservation when partial usage is unknown, or reconcile it from collected stream usage instead of treating the request as unused.
There was a problem hiding this comment.
This method is the last-resort branch: it only runs when neither the success nor failure logging callback ever fired for the request, which per the existing docstring on _arelease_max_parallel_requests_on_disconnect is distinct from the case where a client reads most of a stream then disconnects (that's the "partial-spend billing" event, which fires its own success callback with real usage and reconciles normally through the existing path, never reaching this method). In the actual case this method handles, there is no partial usage available to reconcile against at all, so refunding the full reservation is strictly better than the pre-fix behavior of leaking it until the window's TTL. Note the disable_streaming_logging flag cited in the finding doesn't exist in this codebase. Not changing this further; flagging for visibility.
There was a problem hiding this comment.
The fallback remains reachable after delivered chunks: disable_streaming_logging exists and returns False at litellm/proxy/common_request_processing.py:193-194, while assembly failures do likewise at lines 217-221; cleanup then invokes the full-refund path at lines 2713-2720. That path reconciles both reservations with actual_tokens=0 at litellm/proxy/hooks/parallel_request_limiter_v3.py:4123-4140, so consumed input/output tokens can remain uncharged.
PR overviewThis pull request adds project-level per-model input and output token-per-minute limits to the proxy and integrates them with project configuration and parallel request limiting. Six issues remain open, including a project configuration handling flaw that can disable configured model RPM, TPM, ITPM, and OTPM enforcement. Several request and streaming edge cases also allow project keys to under-reserve tokens or refund consumed capacity, enabling quota bypass and excess provider usage. Three issues have already been addressed, but the remaining enforcement gaps should be resolved before relying on these limits as a security or cost-control boundary. Open issues (6)
Fixed/addressed: 3 · PR risk: 6/10 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
… image content
Two remaining spots still used data.get("input") is not None alone to
detect embeddings for the project TPM/OTPM output-cap logic, missing the
call_type check the estimator itself already uses. That misclassified
Responses API calls as embeddings and skipped the max_tokens output cap
entirely, letting an unbounded generation blow past OTPM before post-call
reconciliation.
Also route Responses API input through the standard
transform_responses_api_input_to_messages helper before token counting.
token_counter's text argument only joins plain strings in a list, so an
input_image content block contributed close to zero tokens to the ITPM
estimate instead of the real image token count.
|
@greptile-apps review |
- _refund_reserved_tokens no-op when amount is zero - reserve_io_tokens early-return when no ITPM/OTPM descriptors present - reserve_io_tokens ITPM-only path (no OTPM descriptors) - _warn_project_io_token_and_tpm_coexist_once suppression after first warn - _strip_audio_content_blocks passthrough for non-list, non-dict, and non-list-content edge cases - OTPM over-limit path releases a stashed parallel slot before raising - ITPM/OTPM-only request (no combined TPM/RPM) stores rate-limit response status in data - _get_reserved_itpm/otpm_tokens_from_kwargs returns 0 for non-numeric corrupted stash - _narrow_reserved_scopes returns empty set for non-list input - _resolve_io_token_reconcile_usage: ResponseAPIUsage with cached tokens, dict-shaped usage, and unrecognised usage type - _build_io_token_reservation_ops returns empty list when usage is unresolvable - disconnect cleanup returns early when reservation already released - async_post_call_failure_hook skips RPM-only descriptors (no tokens_per_unit) in the combined-TPM refund loop - exception in disconnect IO refund is swallowed, not propagated Also removes one dead line: the after _handle_rate_limit_error in _reserve_project_io_tokens_or_raise (unreachable since that function always raises).
|
@greptileai review |
Previously every input_audio block reserved a flat DEFAULT_AUDIO_TOKEN_ESTIMATE (300) tokens regardless of audio length. A caller could pack a 10-minute recording into a single block and issue concurrent requests that each reserved only 300 tokens, exhausting the project ITPM quota while paying negligible quota cost. _count_audio_content_blocks is replaced by two new methods: - _estimate_audio_block_tokens: derives the token estimate for one block from the base64 payload length (decoded_bytes // _AUDIO_BYTES_PER_TOKEN), assuming the lowest reasonable bitrate (8 kHz mono PCM-16, 1600 bytes/token) so longer recordings always reserve proportionally more. Floored at DEFAULT_AUDIO_TOKEN_ESTIMATE (so reference-only blocks and genuinely short clips still get a non-trivial reservation) and capped at _MAX_AUDIO_TOKENS_PER_BLOCK (6000, roughly 10 minutes) so malformed or enormous payloads can't claim unbounded reservations. - _estimate_audio_content_tokens: sums _estimate_audio_block_tokens across all messages; drives both the strip-before-counting decision and the audio token add-on in _estimate_precise_input_tokens. Also covers remaining uncovered lines on the PR patch with 19 new unit tests (see previous commit), and removes one dead unreachable return statement. Co-authored-by: Cursor <cursoragent@cursor.com>
| token_counter( | ||
| model=model or "", | ||
| messages=countable_messages, | ||
| text=None if countable_messages is not None else (data.get("prompt") or data.get("input")), |
There was a problem hiding this comment.
Medium: Referenced files bypass ITPM reservation
The Responses transformer converts input_file into a file content block, which token_counter rejects as unsupported. The exception fallback then counts only the serialized file_id or URL, not the referenced document, so a project key can concurrently submit a large uploaded file while each request reserves only a handful of ITPM tokens. Resolve referenced file size/token usage before reservation, or apply a conservative reservation and reject file references when their content cannot be measured.
There was a problem hiding this comment.
A complete fix requires fetching file content from the Files API at pre-call time, which adds a mandatory network round-trip to every Responses request. The post-call reconciliation path already adjusts the ITPM counter to actual usage after the call, so the window is bounded to the in-flight period. Deferring to a follow-up.
…adata project_metadata was populated from _project_obj.metadata (the freeform JSONB column), so get_project_model_itpm_limit / get_project_model_otpm_limit always returned None even when the dedicated model_itpm_limit / model_otpm_limit columns were set. Add LiteLLM_ProjectTable.merged_metadata, which overlays the four dedicated rate-limit columns on top of the base metadata dict (None columns are skipped so legacy metadata values are not clobbered). Switch all four call-sites in user_api_key_auth.py to use merged_metadata instead of .metadata.
|
@greptileai review |
|
|
||
| messages = data.get("messages") | ||
| if messages is None and call_type in RESPONSES_API_CALL_TYPES and data.get("input") is not None: | ||
| messages = self._responses_input_to_chat_messages(data) |
There was a problem hiding this comment.
Low: Chained Responses requests bypass ITPM reservation
This transforms only the current input; the Responses bridge appends history referenced by previous_response_id later in async_responses_api_session_handler. A project key can create a large response and launch concurrent tiny continuations referencing it, with each request reserving only the tiny new input while the provider processes the full chain. Resolve and count the stored context before reservation, or conservatively reserve its recorded input usage.
There was a problem hiding this comment.
This is a pre-existing limitation shared with model_tpm_limit. Fixing it properly requires fetching stored context token counts from the spend log before reservation , a synchronous DB call in the pre-call hot path for all Responses API traffic.
|
@greptileai review |
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
…/TPM Both has_explicit_max_tokens checks (TPM path and OTPM path) only looked at max_tokens and max_completion_tokens, so a Responses API request that set max_output_tokens was treated as unbounded. The hook then injected a smaller data["max_tokens"] = capped_floor, silently truncating the response despite the reservation already honoring the caller's max_output_tokens. Add max_output_tokens to both checks so Responses API requests with an explicit output bound are excluded from the cap-injection branch. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@greptileai review |
1 similar comment
|
@greptileai review |
| "model_rpm_limit": self.model_rpm_limit, | ||
| "model_tpm_limit": self.model_tpm_limit, | ||
| }.items() | ||
| if v is not None |
There was a problem hiding this comment.
Medium: Project rate-limit bypass
The project request validators move these limits into metadata, but the Prisma columns are non-null and default to {}. Consequently, every loaded row includes non-None empty dictionaries here, which overwrite the configured metadata limits before authentication passes them to the limiter; a project key can then exceed its configured model RPM, TPM, ITPM, and OTPM quotas. Persist these fields in the dedicated columns, or merge them in a way that preserves metadata-backed limits while providing an explicit migration/clearing strategy.
…ens for Responses calls The Responses-to-chat-completion transformation only reads max_output_tokens, so writing the implicit output cap to data["max_tokens"] was silently dropped before provider dispatch, letting an unbounded Responses generation blow past the TPM/OTPM budget. Both injection sites now go through one shared helper that picks the field the request type actually honors.
|
@greptileai review |
…eservation explicit_max_tokens was resolved via a truthy `or` chain across max_tokens/max_completion_tokens/max_output_tokens, so an explicit 0 in a non-last field was folded away as absent. Separately, the OTPM reservation unconditionally floored estimated_output_tokens to at least 1, overriding even a correctly-resolved explicit zero and false-rejecting a zero-output request against an already-exhausted OTPM bucket. Both are fixed: the resolver now returns the first non-None field regardless of falsiness, and the OTPM floor is skipped whenever the caller set an explicit output cap.
| a legitimate zero-output request -- counts as explicit. | ||
| """ | ||
| return ( | ||
| data.get("max_tokens") is not None |
There was a problem hiding this comment.
Low: Ignored output fields bypass token reservation
These checks need to be endpoint-specific. For example, /v1/responses accepts raw JSON here but its bridge only applies max_output_tokens; sending max_tokens: 1 therefore reserves one OTPM token and suppresses the implicit cap, while the field is later ignored and the model can generate an unbounded response. The inverse applies to chat requests carrying max_output_tokens. Select the recognized cap field from call_type and ignore unsupported aliases when estimating or deciding whether to inject a cap.
|
@greptileai review |
| if not self._has_explicit_output_cap(data): | ||
| estimated_output_tokens = max(estimated_output_tokens, 1) |
There was a problem hiding this comment.
When an OpenAI Responses request sets max_output_tokens=0, this branch skips the output reservation floor and reserves zero tokens, but the downstream OpenAI transformation raises the effective cap to 16. The provider can therefore generate up to 16 unreserved output tokens, allowing concurrent requests to exceed the project OTPM limit before reconciliation.
Knowledge Base Used: Proxy Server Request Flow
| """ | ||
| from litellm import token_counter | ||
|
|
||
| messages = data.get("messages") |
There was a problem hiding this comment.
Low: Ignored fields shadow billable input
/v1/responses passes raw JSON through this hook before unsupported parameters are filtered, so a caller can send a large input together with messages: []. This branch counts the empty messages, while the Responses transformation later discards that field and sends the large input, allowing concurrent requests to exceed the project's ITPM limit. Select the token-counting source by call_type; Responses calls should always transform input, while embeddings should count input, completions should count prompt, and chat calls should count messages.
1 participant
We already support model_tpm_limit/model_rpm_limit per project, but that treats input and output tokens as one combined bucket, which doesn't match how Bedrock Mantle actually meters usage. This adds model_itpm_limit/model_otpm_limit so a project can cap input and output tokens separately, same shape as the existing fields.
There's already deployment-level itpm/otpm support in router_utils/pre_call_checks/io_token_rate_limit_check.py, but that code only ever deals with one deployment at a time, whereas the proxy hook (parallel_request_limiter_v3.py) has to juggle key/team/org/project limits together on the same request. Rather than bolt a second copy of the reservation logic on, this just adds project-scoped ITPM/OTPM as two more descriptors into the same reservation path the hook already uses for TPM/RPM.
If a project sets both the combined TPM limit and the new itpm/otpm limits on the same model, both get enforced, not one overriding the other, matching what the deployment-level check already does. Input and output are reserved separately up front, and if one side reserves fine but the other is over limit, we roll back the side that already went through so nothing gets left over-counted. On the way back, cached prompt tokens get excluded from the ITPM count since that's how Bedrock bills it, but that's purely for the rate limit math, cost and usage logging still see the full token count.
Problem this solves:
How it solves it:
Relevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Type
🆕 New Feature
Changes
Adds model_itpm_limit/model_otpm_limit to LiteLLM_ProjectTable, NewProjectRequest, and UpdateProjectRequest, routed into project metadata the same way model_tpm_limit already is, with a matching Prisma migration.
Extends parallel_request_limiter_v3.py with project-scoped ITPM/OTPM rate-limit descriptors, built alongside the existing combined-TPM descriptor rather than replacing it, so both are enforced together when a project configures both (with a one-time warning), matching how the deployment-level itpm/otpm check already handles that overlap. Input and output token estimates are reserved atomically per bucket ahead of the call, with the input-side reservation rolled back if the output side turns out to be over limit. On success, the reservation reconciles against actual usage, excluding cached prompt tokens from the input count to match Bedrock's own metering, without touching cost or usage logging elsewhere. On failure, both reservations are refunded.
Adds get_project_model_itpm_limit/get_project_model_otpm_limit to auth_utils.py for parity with the existing key/team/project rate-limit accessors.
Fixes review findings from Greptile and veria-ai: two spots (the combined-TPM and OTPM output-cap checks in parallel_request_limiter_v3.py) still classified any request with data["input"] set as an embedding, which also misclassifies the Responses API (it puts its prompt in "input" too, but does generate output) and skipped the output cap entirely for it. Both now check call_type the same way the token estimator already does. Also routes Responses API input through the standard transform_responses_api_input_to_messages helper before token counting, since token_counter's text argument only joins plain strings in a list and was silently dropping input_image content blocks from the ITPM estimate.
Fixes a second Greptile finding on the same code: both output-cap spots wrote the implicit cap to data["max_tokens"], but the Responses-to-chat-completion transformation only reads max_output_tokens, so the cap was silently dropped before provider dispatch for Responses calls, letting an unbounded generation blow past the reserved TPM/OTPM budget. Both call sites now go through one shared helper (_apply_implicit_output_cap) that picks the field the request type actually honors.
Fixes a third Greptile finding: explicit_max_tokens was resolved via a truthy max_tokens or max_completion_tokens or max_output_tokens chain, so an explicit 0 in a non-last field got folded away as absent. The project OTPM reservation also unconditionally floored estimated_output_tokens to at least 1, which overrode even a correctly-resolved explicit zero and could false-reject a genuine zero-output request against an already-exhausted OTPM bucket. The resolver now returns the first non-None field regardless of falsiness, and the OTPM floor is skipped whenever the caller set an explicit output cap.
Out of scope for this PR: key/team/org-level model_itpm_limit/model_otpm_limit, and Admin UI support (only the auto-generated schema.d.ts types are updated; no form fields yet).
Final Attestation
Docs PR BerriAI/litellm-docs#638