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 # ---------------------------------------------------------------------------