feat(proxy): enqueued-token rate limiting for batches with refund on completion and cancellation - #37539
Merged
mateo-berri merged 9 commits intoAug 20, 2026
Conversation
…completion and cancellation
Contributor
Greptile SummaryThe PR adds an opt-in outstanding-token allowance for batch submissions, including reservation, terminal-state refund, and administrator-only configuration
Confidence Score: 4/5The production change appears safe, but the live Redis test must be moved to the end-to-end suite before merging The previously reported test isolation failure remains: the unit-test module probes localhost and executes against a real Redis server Files Needing Attention: tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
|
| Filename | Overview |
|---|---|
| litellm/proxy/hooks/batch_enqueued_tokens.py | Adds Redis and in-memory accounting for outstanding batch-token reservations, including idempotent refund records |
| litellm/proxy/hooks/batch_rate_limiter.py | Routes opted-in batch requests through the new enqueued-token allowance instead of ordinary per-minute charging |
| litellm/proxy/hooks/parallel_request_limiter_v3.py | Persists successful batch reservations and refunds them on failure, cancellation, or terminal responses |
| litellm/proxy/auth/auth_utils.py | Adds the shared authorization guard that prevents non-admins from changing the allowance |
| litellm/proxy/management_endpoints/key_management_endpoints.py | Applies the administrator-only metadata guard across key creation, update, bulk update, and regeneration paths |
| litellm/proxy/management_endpoints/team_endpoints.py | Applies the administrator-only metadata guard to team creation and updates |
| tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py | Covers the accounting lifecycle but still places a live Redis integration test in the network-free unit-test tree |
Reviews (9): Last reviewed commit: "Merge remote-tracking branch 'origin/lit..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" | ||
|
|
||
| BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset( | ||
| {"completed", "complete", "failed", "expired", "cancelled", "cancelling"} |
Contributor
There was a problem hiding this comment.
Medium: Nonterminal cancellations refund reservations
cancelling represents an ongoing provider transition, but including it here immediately releases the entire reservation. A caller can submit a batch, request cancellation until it returns cancelling, and submit another batch while the first remains nonterminal, repeatedly exceeding the configured enqueued-token allowance. Refund only after the provider reports an actual terminal state.
Suggested change
| {"completed", "complete", "failed", "expired", "cancelled", "cancelling"} | |
| {"completed", "complete", "failed", "expired", "cancelled"} |
Contributor
Author
There was a problem hiding this comment.
Immediate refund on cancellation is this feature's acceptance criterion; the cancelling-to-cancelled window is minutes, and the popped reservation record already prevents double refunds
Contributor
PR overviewThis pull request adds enqueued-token rate limiting for proxy batch requests, reserving capacity when batches are submitted and refunding it when they complete or are cancelled. It also integrates reservation refunds into batch request failure handling. Five issues have been addressed, but two refund paths can still release token reservations while provider-side batches may remain active. Callers could exploit cancellation transitions or post-dispatch timeouts to submit concurrent work beyond the configured enqueued-token allowance, weakening the intended resource-control boundary. Refunds should be limited to confirmed terminal states or requests known not to have been dispatched. Open issues (2)
Fixed/addressed: 5 · PR risk: 6/10 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…s Cluster compatibility
Contributor
Author
… by backend, lowercase terminal statuses Reserve-script failures now roll back the scopes already incremented before re-raising into the in-memory fallback, so a partial redis outage no longer leaks counter increments that shrink the shared allowance. Reservations record which backend granted them, so a refund never debits redis counters an in-memory grant did not charge. Terminal-status matching is now case-insensitive because the Bedrock async-invoke retrieve path returns raw AWS-cased statuses like Completed.
Contributor
Author
…orker In-memory grants now record an owner token, and a refund only debits local counters when the popping worker is the one that granted them, so a terminal response handled elsewhere can no longer shrink another worker's unrelated fallback reservations. A Redis-granted refund that fails no longer falls back to decrementing local counters either: the leaked Redis increments expire with the TTL and only tighten the allowance.
Contributor
Author
Contributor
Author
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.
Autofix Details
Bugbot Autofix prepared fixes for both issues found in the latest run.
- ✅ Fixed: Over-limit Redis errors grant quota
- Routed the mid-reserve rollback through the best-effort
_rollback_partial_reserveso a failing DECRBY on prior scopes no longer bubbles out of_reserve_via_redisand drops through to_reserve_in_memory, preserving the intendedBatchEnqueuedTokenOverLimitresponse.
- Routed the mid-reserve rollback through the best-effort
- ✅ Fixed: In-memory reservation refunds are skipped
pop_reservationnow falls back to_pop_local_recordwhenever the Redis pop returnsNone(not just on Redis exceptions), so cancel/complete refunds recover reservations thatsave_reservationhad to stash locally when Redis was unavailable.
Or push these changes by commenting:
@cursor push 3c4eecac5b
Preview (3c4eecac5b)
diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py
--- a/litellm/proxy/hooks/batch_enqueued_tokens.py
+++ b/litellm/proxy/hooks/batch_enqueued_tokens.py
@@ -246,7 +246,7 @@ async def _reserve_via_redis(
already_reserved=scopes[:index],
)
if result[0] != 1:
- await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes[:index])
+ await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index])
return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1])
return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="redis")
@@ -246,7 +246,7 @@ async def _reserve_via_redis(
already_reserved=scopes[:index],
)
if result[0] != 1:
- await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes[:index])
+ await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index])
return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1])
return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="redis")
@@ -384,8 +384,7 @@ async def pop_reservation(
verbose_proxy_logger.warning(
"Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e)
)
- raw = await self._pop_local_record(batch_id, litellm_parent_otel_span)
- else:
+ if raw is None:
raw = await self._pop_local_record(batch_id, litellm_parent_otel_span)
if raw is None:
return None
@@ -384,8 +384,7 @@ async def pop_reservation(
verbose_proxy_logger.warning(
"Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e)
)
- raw = await self._pop_local_record(batch_id, litellm_parent_otel_span)
- else:
+ if raw is None:
raw = await self._pop_local_record(batch_id, litellm_parent_otel_span)
if raw is None:
return None
diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
--- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
+++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
@@ -132,14 +132,25 @@ def __init__(
self,
fail_reserve_keys: frozenset[str] = frozenset(),
fail_refund_keys: frozenset[str] = frozenset(),
+ fail_save_keys: frozenset[str] = frozenset(),
) -> None:
self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = ()
self.counters: Mapping[str, int] = MappingProxyType({})
+ self.records: Mapping[str, str] = MappingProxyType({})
self.fail_reserve_keys = fail_reserve_keys
self.fail_refund_keys = fail_refund_keys
+ self.fail_save_keys = fail_save_keys
def async_register_script(self, script: str):
- kind: Final = "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "record"
+ kind: Final = (
+ "reserve"
+ if "INCRBY" in script
+ else "refund"
+ if "DECRBY" in script
+ else "save"
+ if "SET" in script
+ else "pop"
+ )
async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object:
self.script_calls = (*self.script_calls, (kind, tuple(keys)))
@@ -132,14 +132,25 @@ def __init__(
self,
fail_reserve_keys: frozenset[str] = frozenset(),
fail_refund_keys: frozenset[str] = frozenset(),
+ fail_save_keys: frozenset[str] = frozenset(),
) -> None:
self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = ()
self.counters: Mapping[str, int] = MappingProxyType({})
+ self.records: Mapping[str, str] = MappingProxyType({})
self.fail_reserve_keys = fail_reserve_keys
self.fail_refund_keys = fail_refund_keys
+ self.fail_save_keys = fail_save_keys
def async_register_script(self, script: str):
- kind: Final = "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "record"
+ kind: Final = (
+ "reserve"
+ if "INCRBY" in script
+ else "refund"
+ if "DECRBY" in script
+ else "save"
+ if "SET" in script
+ else "pop"
+ )
async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object:
self.script_calls = (*self.script_calls, (kind, tuple(keys)))
@@ -167,6 +178,15 @@ def _run(self, kind: str, keys: tuple[str, ...], args: tuple[str | bytes | int |
else {**self.counters, keys[0]: remaining}
)
return 1
+ if kind == "save":
+ if keys[0] in self.fail_save_keys:
+ raise ConnectionError(f"simulated redis failure for {keys[0]}")
+ self.records = MappingProxyType({**self.records, keys[0]: str(args[0])})
+ return 1
+ if kind == "pop":
+ value: Final = self.records.get(keys[0])
+ self.records = MappingProxyType({key: val for key, val in self.records.items() if key != keys[0]})
+ return value
raise AssertionError(f"unexpected {kind} script call for keys {keys}")
@@ -167,6 +178,15 @@ def _run(self, kind: str, keys: tuple[str, ...], args: tuple[str | bytes | int |
else {**self.counters, keys[0]: remaining}
)
return 1
+ if kind == "save":
+ if keys[0] in self.fail_save_keys:
+ raise ConnectionError(f"simulated redis failure for {keys[0]}")
+ self.records = MappingProxyType({**self.records, keys[0]: str(args[0])})
+ return 1
+ if kind == "pop":
+ value: Final = self.records.get(keys[0])
+ self.records = MappingProxyType({key: val for key, val in self.records.items() if key != keys[0]})
+ return value
raise AssertionError(f"unexpected {kind} script call for keys {keys}")
@@ -252,6 +272,39 @@ async def test_failed_redis_refund_leaves_local_counters_untouched():
assert fake.counters == {counter_key: 60}
+@pytest.mark.asyncio
+async def test_over_limit_still_rejects_when_prior_scope_rollback_fails():
+ key_scope = _scope(limit=100, key="api_key")
+ team_scope = _scope(limit=50, key="team")
+ fake = _SingleKeyRedisFake(fail_refund_keys=frozenset({f"batch_enqueued_tokens:api_key:{key_scope.value}"}))
+ store = BatchEnqueuedTokenStore(
+ internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
+ )
+ seed = await store.reserve(tokens=30, scopes=(key_scope, team_scope))
+ assert isinstance(seed, BatchEnqueuedTokenReservation)
+
+ outcome = await store.reserve(tokens=25, scopes=(key_scope, team_scope))
+ assert outcome == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=30)
+
+
+@pytest.mark.asyncio
+async def test_pop_reservation_falls_back_to_local_when_save_wrote_locally():
+ key_scope = _scope(limit=100, key="api_key")
+ record_key: Final = "batch_enqueued_token_reservation:batch_local"
+ fake = _SingleKeyRedisFake(fail_save_keys=frozenset({record_key}))
+ store = BatchEnqueuedTokenStore(
+ internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
+ )
+ reservation = BatchEnqueuedTokenReservation(tokens=40, scopes=(key_scope,), backend="redis")
+
+ await store.save_reservation("batch_local", reservation)
+ assert record_key not in fake.records
+
+ popped = await store.pop_reservation("batch_local")
+ assert popped == reservation
+ assert await store.pop_reservation("batch_local") is None
+
+
@pytest.mark.asyncio
async def test_pop_reservation_defaults_legacy_records_to_redis_backend():
store = _in_memory_store()
@@ -252,6 +272,39 @@ async def test_failed_redis_refund_leaves_local_counters_untouched():
assert fake.counters == {counter_key: 60}
+@pytest.mark.asyncio
+async def test_over_limit_still_rejects_when_prior_scope_rollback_fails():
+ key_scope = _scope(limit=100, key="api_key")
+ team_scope = _scope(limit=50, key="team")
+ fake = _SingleKeyRedisFake(fail_refund_keys=frozenset({f"batch_enqueued_tokens:api_key:{key_scope.value}"}))
+ store = BatchEnqueuedTokenStore(
+ internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
+ )
+ seed = await store.reserve(tokens=30, scopes=(key_scope, team_scope))
+ assert isinstance(seed, BatchEnqueuedTokenReservation)
+
+ outcome = await store.reserve(tokens=25, scopes=(key_scope, team_scope))
+ assert outcome == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=30)
+
+
+@pytest.mark.asyncio
+async def test_pop_reservation_falls_back_to_local_when_save_wrote_locally():
+ key_scope = _scope(limit=100, key="api_key")
+ record_key: Final = "batch_enqueued_token_reservation:batch_local"
+ fake = _SingleKeyRedisFake(fail_save_keys=frozenset({record_key}))
+ store = BatchEnqueuedTokenStore(
+ internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
+ )
+ reservation = BatchEnqueuedTokenReservation(tokens=40, scopes=(key_scope,), backend="redis")
+
+ await store.save_reservation("batch_local", reservation)
+ assert record_key not in fake.records
+
+ popped = await store.pop_reservation("batch_local")
+ assert popped == reservation
+ assert await store.pop_reservation("batch_local") is None
+
+
@pytest.mark.asyncio
async def test_pop_reservation_defaults_legacy_records_to_redis_backend():
store = _in_memory_store()You can send follow-ups to the cloud agent here.
…ailure, find locally saved records on pop A Redis over-limit verdict now survives a failing rollback DECRBY instead of escaping into the in-memory fallback and granting tokens the counter already rejected; the unrolled increments expire with the TTL. pop_reservation now falls through to the local record when the Redis pop succeeds but finds nothing, so a reservation saved in memory after a transient Redis save failure still refunds on cancel or completion.
Contributor
Author
…oxy admins The field replaces the standard RPM/TPM checks for batch submissions, so a key holder or team admin writing it could pick their own batch quota. Mirrors the output-token-estimate admin gate: change-based, so resending the stored value stays allowed, and enforced on key generate, update, bulk team-key update, regenerate, and team new/update.
|
hiraku-miyoshi seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Contributor
Author
Contributor
Author
|
bugbot run |
Contributor
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Local fallback can double-refund quota
- Eliminated the dual-write at its source: save_reservation now returns without writing a local copy when the Redis save raises, since a Redis SET can still land server-side after the client times out and a local fallback would let a later pop refund again.
Or push these changes by commenting:
@cursor push 2b0b853a67
Preview (2b0b853a67)
diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py
--- a/litellm/proxy/hooks/batch_enqueued_tokens.py
+++ b/litellm/proxy/hooks/batch_enqueued_tokens.py
@@ -355,12 +355,11 @@ async def save_reservation(
(self._record_key(batch_id),),
(serialized, BATCH_ENQUEUED_TOKEN_TTL_SECONDS),
)
- except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record
+ except Exception as e: # noqa: BLE001 # a Redis SET can land server-side after the client raises; a local fallback would risk a double refund via pop_reservation
verbose_proxy_logger.warning(
- "Redis enqueued-token reservation save failed, falling back to in-memory: %s", str(e)
+ "Redis enqueued-token reservation save failed; refund may leak until TTL: %s", str(e)
)
- else:
- return
+ return
await self.internal_usage_cache.async_set_cache(
key=self._record_key(batch_id),
value=serialized,
@@ -355,12 +355,11 @@ async def save_reservation(
(self._record_key(batch_id),),
(serialized, BATCH_ENQUEUED_TOKEN_TTL_SECONDS),
)
- except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record
+ except Exception as e: # noqa: BLE001 # a Redis SET can land server-side after the client raises; a local fallback would risk a double refund via pop_reservation
verbose_proxy_logger.warning(
- "Redis enqueued-token reservation save failed, falling back to in-memory: %s", str(e)
+ "Redis enqueued-token reservation save failed; refund may leak until TTL: %s", str(e)
)
- else:
- return
+ return
await self.internal_usage_cache.async_set_cache(
key=self._record_key(batch_id),
value=serialized,
diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
--- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
+++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py
@@ -246,7 +246,10 @@ async def test_over_limit_verdict_survives_a_failing_rollback():
@pytest.mark.asyncio
-async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing():
+async def test_save_reservation_does_not_dual_write_when_redis_save_raises():
+ """LIT-5273 regression: a Redis SET timeout can still land server-side, so writing a local
+ fallback on save failure risks a double refund the first time pop_reservation runs against
+ Redis and a later cancel/retrieve pops the leftover local record."""
scope = _scope(limit=100)
record_key: Final = "batch_enqueued_token_reservation:batch_local_record"
fake = _SingleKeyRedisFake(fail_save_keys=frozenset({record_key}))
@@ -246,7 +246,10 @@ async def test_over_limit_verdict_survives_a_failing_rollback():
@pytest.mark.asyncio
-async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing():
+async def test_save_reservation_does_not_dual_write_when_redis_save_raises():
+ """LIT-5273 regression: a Redis SET timeout can still land server-side, so writing a local
+ fallback on save failure risks a double refund the first time pop_reservation runs against
+ Redis and a later cancel/retrieve pops the leftover local record."""
scope = _scope(limit=100)
record_key: Final = "batch_enqueued_token_reservation:batch_local_record"
fake = _SingleKeyRedisFake(fail_save_keys=frozenset({record_key}))
@@ -258,11 +261,9 @@ async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing():
assert isinstance(reservation, BatchEnqueuedTokenReservation)
await store.save_reservation("batch_local_record", reservation)
assert not fake.records
-
- popped = await store.pop_reservation("batch_local_record")
- assert popped == reservation
- await store.refund(popped)
- assert not fake.counters
+ assert (
+ store.internal_usage_cache.dual_cache.in_memory_cache.get_cache(key=record_key) is None
+ )
assert await store.pop_reservation("batch_local_record") is None
@@ -258,11 +261,9 @@ async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing():
assert isinstance(reservation, BatchEnqueuedTokenReservation)
await store.save_reservation("batch_local_record", reservation)
assert not fake.records
-
- popped = await store.pop_reservation("batch_local_record")
- assert popped == reservation
- await store.refund(popped)
- assert not fake.counters
+ assert (
+ store.internal_usage_cache.dual_cache.in_memory_cache.get_cache(key=record_key) is None
+ )
assert await store.pop_reservation("batch_local_record") is NoneYou can send follow-ups to the cloud agent here.
…s so local ghosts cannot double-refund
Contributor
Author
Contributor
Author
|
bugbot run |
Contributor
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 5ab20c3. Configure here.
Contributor
…ve their counters
Contributor
Author
Contributor
Author
|
bugbot run |
| ) | ||
| stash.parallel_slot = None | ||
|
|
||
| if stash.batch_enqueued_reservation is not None: |
Contributor
There was a problem hiding this comment.
Low: Provider-side timeouts refund accepted batches
This refunds the reservation for every failure, although the batch endpoint also invokes this hook for exceptions after provider dispatch. Because the batch request accepts a caller-controlled timeout, a user can tune repeated requests to time out after the provider accepts them, restoring the allowance while those upstream batches continue processing. Only refund when the request was never dispatched; for an ambiguous provider outcome, retain the reservation until the batch is reconciled or its TTL expires.
Contributor
Author
There was a problem hiding this comment.
Ambiguous timeouts leave no batch id to reconcile against; retaining would lock the allowance 8 days on every transient failure. Documented as bounded residual risk.
Contributor
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 2a771ca. Configure here.
…itellm_batch_enqueued_token_limit
Contributor
Author
Contributor
Author
|
bugbot run |
Contributor
Author
|
bugbot run |
Contributor
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit d5ac495. Configure here.
mateo-berri
enabled auto-merge
August 20, 2026 04:11
yucheng-berri
approved these changes
Aug 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TLDR
Problem this solves:
disable_batch_input_file_rate_limiting, drops all batch quota controlHow it solves it:
batch_enqueued_token_limitin key or team metadataUser Flow
Before: a team whose batch jobs outsize their per-minute limits cannot submit batches at all without turning batch rate limiting off entirely
{"rpm_limit": 2, "tpm_limit": 10}purpose=batch, getting back afile-...id{"input_file_id": "file-..."}disable_batch_input_file_rate_limiting: true, after which any key can enqueue unlimited batch workAfter: the same team's key carries an enqueued-token allowance, so batch submission is governed by outstanding batch tokens instead of per-minute limits
{"rpm_limit": 2, "tpm_limit": 10, "metadata": {"batch_enqueued_token_limit": 100000}}purpose=batch, getting back afile-...id{"input_file_id": "file-..."}validating, even though 3 rows exceed the key's 2 RPMRelevant issues
Linear ticket
Resolves LIT-5273
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
Both runs hit a live proxy booted from a fresh worktree with a fresh Postgres DB, local Redis, and real OpenAI batches on gpt-5.4-mini (real $$). Before = merge base a1afc2f, After = PR tip d5ac495, same four cases each side.
batch_input.jsonlis 5 tiny gpt-5.4-mini chat rows withmax_completion_tokens: 16Before (a1afc2f, merge base)
Case 1: batch over the key's RPM with an allowance in metadata, 429 and the allowance is ignored
Case 2: 1-token allowance and no per-minute limits, the batch sails through because nothing reads the allowance
Case 3: allowance sized to fit one batch, then cancel, every create is 200 because nothing is accounted
Case 4: online chat on a 2 RPM key with the allowance present, 200/200/429
After (d5ac495, PR tip)
Case 1: the same over-RPM batch with the 100k allowance is now accepted
Case 2: the 1-token allowance now rejects at the gateway, before the provider sees anything
The 429 names the reservation size N = 65 tokens, so the refund case sizes its allowance at 65 + 65//2 = 97
Case 3: 97-token allowance, exhaustion 429, cancel refunds, retry succeeds
Redis on 6379 held the live
batch_enqueued_tokens:api_key:be176bd8...counter at 65 after the refund (not 130), batch 1's popped reservation key held the empty-string tombstone (EXISTS=1, STRLEN=0) that blocks a second refund, and batch 3's reservation showed{"tokens":65,...,"backend":"redis","reserved_at_monotonic":176584.85...}with record and tombstone TTLs (691173, 691174) at or below the counter's (691174), so a reservation record can never outlive the counter it refunds; the proxy log shows zero in-memory-fallback lines for enqueued-token opsCase 4: online chat on the same-shaped 2 RPM key, unchanged 200/200/429
Admin gate: a non-admin cannot grant themselves an allowance
Surprises:
Type
🆕 New Feature
✅ Test
Caveats (if any)
batch_enqueued_token_limitis proxy-admin-only across key generate/update/regenerate/bulk update and team new/update; edit forms that resend the stored value unchanged stay allowedLEGACY_MULTI_INSTANCE_RATE_LIMITING=true) has no batch handling, so the feature is a no-op there, matching existing batch limitsQA runbook
{"metadata": {"batch_enqueued_token_limit": 1}}, upload a 3-request JSONL with it, create the batch, and expect 429 naming "Batch enqueued token limit exceeded" plus "Batch requires N tokens"cancellingFinal Attestation
The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
d5ac495 passes /live-pr-risk
Note
Cursor Bugbot is generating a summary for commit d5ac495. Configure here.