scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(rate-limiting): revert refunding the raising key's own ambiguous outcome

Browse source 
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.
This commit is contained in:
Deepanshu 2026-08-18 12:46:22 -04:00
parent 23b2a750fe
commit 66def53981
2 changed files with 57 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/proxy/hooks/tag_rate_limiter.py
View file

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

40
tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py
View file

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