From 66def53981a48710e9e943bd7885dc505465cab5 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 18 Aug 2026 12:46:22 -0400 Subject: [PATCH] fix(rate-limiting): revert refunding the raising key's own ambiguous outcome The previous fix (refund through the raising index inclusive) was itself a regression: these are shared, chain-wide buckets with no per-request ownership tracking, so decrementing a key whose own increment might or might not have committed is just as likely to erase a different, legitimately-admitted concurrent request's charge on the same key as it is to undo our own. An attacker could repeatedly cancel requests to deliberately erase other callers' charges and exceed the configured limit -- worse than the alternative this reverts to: a committed-but-unrefunded key self-heals via its own TTL. Only strictly earlier admissions in the same batch (this request's own confirmed-successful increments, never ambiguous) are refunded now. --- litellm/proxy/hooks/tag_rate_limiter.py | 48 +++++++++++-------- .../proxy/hooks/test_tag_rate_limiter.py | 40 +++++++++++----- 2 files changed, 57 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/hooks/tag_rate_limiter.py b/litellm/proxy/hooks/tag_rate_limiter.py index 927c7bd0981..2bed00989a4 100644 --- a/litellm/proxy/hooks/tag_rate_limiter.py +++ b/litellm/proxy/hooks/tag_rate_limiter.py @@ -833,16 +833,27 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer A later key's own admission raising (a transient Redis error, or this coroutine being cancelled mid-call, e.g. the caller disconnecting) is treated the same as a normal rejection for refund - purposes, with one difference: a clean rejection is guaranteed by - TAG_RL_CHECK_AND_INCR_SCRIPT to never have incremented that key (it - returns before calling INCRBY), so only the earlier admissions need - refunding. A raise gives no such guarantee -- Redis can commit the - INCRBY and still have the call raise if the response back to us is - lost (a timeout, a dropped connection) -- so that key's own possibly - -committed increment is refunded too. Refunding a key that in fact - never committed is harmless (floors at 0); skipping one that did - commit would leak a permanently-charged counter or concurrency - reservation for the rest of that key's TTL. + purposes: only the earlier admissions in this batch are refunded, + never the raising key's own key. This is deliberate, not an + oversight: a raise gives no guarantee that key's own increment + didn't already commit server-side (Redis can run the INCRBY and + still have the call raise if the response back to us is lost), but + these are shared, chain-wide buckets with no per-request ownership + tracking -- decrementing on that guess is just as likely to erase a + *different*, legitimately-admitted concurrent request's charge on + the same key as it is to undo our own. That failure mode (an + attacker repeatedly cancelling requests to erase other callers' + charges and exceed the configured limit) is worse than the + alternative this accepts instead: a key that did commit but never + gets refunded self-heals via its own TTL -- see `_ttl_for`. The + earlier admissions refunded here are never ambiguous like this: they + are this same request's own confirmed-successful increments, so + undoing them is always safe. + + Refunds are best-effort: a refund that fails (e.g. a transient Redis + error) is logged and skipped rather than raised, so one bad refund + can't stop the rest of the batch from being refunded, and can't turn + a clean rejection into an unhandled exception. Returns (failing_index, values). On success, failing_index is None and values holds each key's new post-increment value, same order as @@ -860,19 +871,16 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection; see comment above for index, (cache, key, limit, increment, ttl) in enumerate(checks): admitted = False - completed = False try: admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl) - completed = True finally: - # completed=False means the awaited call itself raised or - # was cancelled -- refund through this index inclusive, per - # the docstring above. completed=True and admitted=False is - # a clean rejection -- refund only the earlier ones, since - # this key's own increment never happened. - if not completed: - await self._refund_admitted(checks, up_to_index=index + 1) - elif not admitted: + # Runs on a normal rejection (admitted stays False) and on + # any exception/cancellation from the awaited call above + # (admitted never gets assigned, so it's still the False set + # just before the try) -- either way, only the earlier, + # known-safe admissions are refunded; see the docstring + # above for why this key's own ambiguous outcome is not. + if not admitted: await self._refund_admitted(checks, up_to_index=index) if admitted: admitted_values.append(value) # mutable-ok: see accumulator comment above diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py index bef93da4752..fbc4a9ba119 100644 --- a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py @@ -2084,17 +2084,22 @@ async def test_exception_mid_batch_refunds_every_earlier_admission_before_propag @pytest.mark.asyncio -async def test_a_committed_increment_whose_response_is_lost_is_also_refunded(time_controller): +async def test_a_raising_keys_own_ambiguous_outcome_is_never_refunded(time_controller): """ - Regression test: a key can commit its own increment (e.g. Redis runs - the INCRBY) and still have the call raise if the response back to us is - lost (a timeout, a dropped connection) -- the caller can't tell a lost - response apart from a call that never reached Redis at all. The earlier - fix only refunded indices *before* the one that raised, leaving this - key's own possibly-committed increment permanently charged. It must be - refunded too, not just the earlier ones in the same batch. + Regression test for a bug this exact fix briefly introduced: a key can + commit its own increment (e.g. Redis runs the INCRBY) and still have + the call raise if the response back to us is lost, so a raise never + proves that key's own attempt didn't commit. But these are shared, + chain-wide buckets with no per-request ownership tracking, so + decrementing on that guess is just as likely to erase a *different*, + legitimately-admitted concurrent request's charge on the same key as it + is to undo our own -- an attacker could repeatedly cancel requests to + erase other callers' charges and exceed the configured limit. The + raising key's own outcome must never be refunded, only strictly earlier + (confirmed-safe) admissions in the same batch. """ - raising_key = "{tag_rl:test:lost-response-refund:a}:requests" + admitted_key = "{tag_rl:test:ambiguous-no-refund:a}:requests" + raising_key = "{tag_rl:test:ambiguous-no-refund:b}:requests" class _FlakyLimiter(_PROXY_TagRateLimiter): async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int): @@ -2109,10 +2114,23 @@ async def test_a_committed_increment_whose_response_is_lost_is_also_refunded(tim flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now) with pytest.raises(RuntimeError): - await flaky._atomic_check_and_increment([(flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60)]) + await flaky._atomic_check_and_increment( + [ + (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60), + ] + ) + # The earlier, confirmed-successful admission in this same batch is + # always safe to refund. + admitted_value = await flaky.internal_usage_cache.async_get_cache(key=admitted_key, litellm_parent_otel_span=None) + assert (float(admitted_value) if admitted_value is not None else 0.0) == 0.0 + + # The raising key's own committed increment must survive -- refunding + # it would be indistinguishable from erasing a different request's + # legitimate charge on the same shared bucket. raising_key_value = await flaky.internal_usage_cache.async_get_cache(key=raising_key, litellm_parent_otel_span=None) - assert (float(raising_key_value) if raising_key_value is not None else 0.0) == 0.0 + assert float(raising_key_value) == 1.0 # ---------------------------------------------------------------------------