mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(rate-limiting): stop fallback hops from double-charging a requests cap
Cursor Bugbot finding on PR #36541, live-confirmed against a real proxy: a "requests" limit is meant to cap logical client requests, not internal routing attempts, but each hop's atomic admission increment had no refund path of its own. A chain that failed once before succeeding burned two units of a one-request-per-period budget for a single logical call; live reproduction showed the retry's own admission rejected with current=1.0 limit=1.0 even though the client only made one request. Generalizes the same next-hop-releases-a-prior-hop's-stale-reservation pattern already used for concurrency: a queued "requests" increment still present when a new hop's admission runs can only belong to an already-failed earlier hop, so it gets refunded there before this hop's own check runs. Unlike concurrency, a hop that goes on to succeed (or is the chain's own final failure) is never refunded, so exactly one unit survives per logical request regardless of how many hops it took. An initial version of this fix also popped the new field in async_log_success_event/async_log_failure_event "for hygiene." That broke live: litellm's has_logged_async_failure dedup lets the first failing hop's own failure event through, not only a chain's final failure, so the pop discarded the entry before the next hop's own admission ever got a chance to refund it, permanently stranding the charge. Both handlers now leave the field completely untouched; a dedicated regression test simulates that exact real failure-event sequence, not just two bare admission calls. Also fixes a related veria-ai finding on the same file: the pending- reservations mirror key embedded litellm_call_id (caller-controlled via the x-litellm-call-id header) directly, with no length bound. Hashed via the same _fixed_length_identity helper already used for tag values.
This commit is contained in:
Deepanshu
parent
a26da3783c
commit
ed726bef7b
2 changed files with 229 additions and 11 deletions
|
|
@ -654,6 +654,28 @@ _CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600
|
|||
# can't be forged or guessed.
|
||||
_PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys"
|
||||
|
||||
# Same `model_call_details`-stashing rationale as the field above, for a
|
||||
# different unit: "requests" is atomic and admitted once per hop (see
|
||||
# _ATOMIC_UNITS), same as concurrency, but a "requests" limit is meant to cap
|
||||
# logical client requests, not internal routing attempts -- a chain that
|
||||
# fails once before succeeding must still consume exactly one unit overall,
|
||||
# not one per hop. _release_stale_hop_reservations refunds a stale entry
|
||||
# here the same way it releases a stale concurrency reservation, since its
|
||||
# own invariant (a queued entry still present when a new hop's admission
|
||||
# runs can only belong to an earlier hop of this same request that already
|
||||
# failed) holds identically for either unit. Unlike concurrency, a
|
||||
# successful (or chain-final-failing) hop's own entry here is deliberately
|
||||
# never refunded -- exactly one unit must survive per logical request -- so
|
||||
# async_log_success_event/async_log_failure_event must leave this field
|
||||
# completely untouched: litellm's has_logged_async_failure dedup lets the
|
||||
# *first* failing hop's own failure event through (not only a chain's final
|
||||
# failure), so popping this field there -- even just to discard it -- would
|
||||
# strand the very entry the *next* hop's admission is relying on being able
|
||||
# to refund. There is no final-hop/cache-mirror problem to solve for this
|
||||
# field either: the one hop that never gets superseded is exactly the one
|
||||
# whose charge should stick, with nothing left to clean up.
|
||||
_PENDING_REQUEST_INCREMENTS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_request_increments"
|
||||
|
||||
# Mirrors the latest hop's own queued reservation in the same external cache
|
||||
# the reservations themselves live in, keyed by (litellm_call_id, key_hash),
|
||||
# for the one release path that cannot reach model_call_details at all:
|
||||
|
|
@ -694,7 +716,10 @@ _PENDING_RESERVATIONS_CACHE_KEY_PREFIX: Final = "model_based_tag_rate_limits:pen
|
|||
|
||||
|
||||
def _pending_reservations_cache_key(call_id: str, key_hash: str | None) -> str:
|
||||
return f"{_PENDING_RESERVATIONS_CACHE_KEY_PREFIX}{call_id}:{key_hash or ''}"
|
||||
# call_id is caller-controlled (the x-litellm-call-id header) with no
|
||||
# length bound -- same unbounded-cache-key concern _fixed_length_identity
|
||||
# documents for tag values, reused here rather than duplicated.
|
||||
return f"{_PENDING_RESERVATIONS_CACHE_KEY_PREFIX}{_fixed_length_identity(call_id)}:{key_hash or ''}"
|
||||
|
||||
|
||||
def _encode_reservations(reservations: Sequence[tuple[str, "_PartitionKey"]]) -> str:
|
||||
|
|
@ -1004,26 +1029,49 @@ def _partition_key(entry: TagRateLimitEntry) -> _PartitionKey:
|
|||
)
|
||||
|
||||
|
||||
def _queue_pending_concurrency_reservations(
|
||||
request_kwargs: Mapping[str, object], reservations: Sequence[tuple[str, _PartitionKey]]
|
||||
def _queue_pending_reservations(
|
||||
request_kwargs: Mapping[str, object], field: str, reservations: Sequence[tuple[str, _PartitionKey]]
|
||||
) -> None:
|
||||
"""Stash reservations on the request's own `model_call_details` -- see
|
||||
`_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring for why this, not a
|
||||
ContextVar or `litellm_call_id`. Silently a no-op without a real logging
|
||||
object (defensive only; every real request has one): the reservation
|
||||
still self-heals via `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, just later.
|
||||
"""Stash reservations on the request's own `model_call_details`, under
|
||||
`field` -- see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring for why this,
|
||||
not a ContextVar or `litellm_call_id`. Silently a no-op without a real
|
||||
logging object (defensive only; every real request has one): a queued
|
||||
concurrency reservation still self-heals via
|
||||
`_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, just later.
|
||||
"""
|
||||
logging_obj: Final = request_kwargs.get("litellm_logging_obj")
|
||||
model_call_details: Final = getattr(logging_obj, "model_call_details", None)
|
||||
if not isinstance(model_call_details, dict):
|
||||
return
|
||||
pending = model_call_details.get(_PENDING_CONCURRENCY_KEYS_FIELD) # rebind-ok: lazily initialized below when absent
|
||||
pending = model_call_details.get(field) # rebind-ok: lazily initialized below when absent
|
||||
if pending is None:
|
||||
pending = [] # mutable-ok: shared, request-scoped accumulator; see field's own docstring # rebind-ok: lazily initialized only when absent
|
||||
model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] = pending
|
||||
model_call_details[field] = pending
|
||||
pending.extend(reservations) # mutable-ok: see comment above
|
||||
|
||||
|
||||
def _pop_reservations(model_call_details: Mapping[str, object], field: str) -> tuple[tuple[str, "_PartitionKey"], ...]:
|
||||
"""No external cache-mirror interaction -- only
|
||||
`_PENDING_CONCURRENCY_KEYS_FIELD` needs that (see its docstring); a
|
||||
"requests" entry here never needs a final-hop release path, so this is
|
||||
the whole mechanism. Snapshots then removes individual items from the
|
||||
same list object rather than a blanket pop of `field` itself, matching
|
||||
`_pop_pending_concurrency_keys`'s own reasoning: a sibling hop sharing
|
||||
this request's `model_call_details` can still be live and appending
|
||||
concurrently, so clearing the whole field here could silently strand
|
||||
that entry instead of it being refunded or left to stand later."""
|
||||
pending = model_call_details.get(field)
|
||||
if not isinstance(pending, list) or not pending:
|
||||
return ()
|
||||
keys: Final = tuple(pending)
|
||||
for key in keys:
|
||||
try:
|
||||
pending.remove(key) # mutable-ok: see queuing helper above
|
||||
except ValueError:
|
||||
pass
|
||||
return keys
|
||||
|
||||
|
||||
def _record_admission_time(request_kwargs: Mapping[str, object], now: float) -> None:
|
||||
"""Stash this hop's admission timestamp -- see `_ADMISSION_TIME_FIELD`'s
|
||||
docstring for why. Silently a no-op without a real logging object
|
||||
|
|
@ -1340,13 +1388,25 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
if configured_limit.unit == "concurrency"
|
||||
)
|
||||
if concurrency_reservations:
|
||||
_queue_pending_concurrency_reservations(resolved_request_kwargs, concurrency_reservations)
|
||||
_queue_pending_reservations(
|
||||
resolved_request_kwargs, _PENDING_CONCURRENCY_KEYS_FIELD, concurrency_reservations
|
||||
)
|
||||
await self._mirror_pending_reservations(
|
||||
resolved_request_kwargs.get("litellm_call_id"),
|
||||
_extract_key_hash(resolved_request_kwargs, metadata_variable_name),
|
||||
concurrency_reservations,
|
||||
)
|
||||
|
||||
request_increments: Final = tuple(
|
||||
(key, _partition_key(configured_limit.entry))
|
||||
for configured_limit, _tag_value, key in atomic_checks
|
||||
if configured_limit.unit == "requests"
|
||||
)
|
||||
if request_increments:
|
||||
_queue_pending_reservations(
|
||||
resolved_request_kwargs, _PENDING_REQUEST_INCREMENTS_FIELD, request_increments
|
||||
)
|
||||
|
||||
return healthy_deployments
|
||||
|
||||
async def _mirror_pending_reservations(
|
||||
|
|
@ -1533,6 +1593,13 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
whose own failure exhausts the retry chain -- async_post_call_failure_hook
|
||||
closes that residual case instead, via the cache mirror
|
||||
`_PENDING_RESERVATIONS_CACHE_KEY_PREFIX` documents.
|
||||
|
||||
The identical invariant -- a queued entry still present here can only
|
||||
belong to an already-failed earlier hop -- holds for a "requests"
|
||||
atomic increment too, so this also refunds any stale entry queued
|
||||
under `_PENDING_REQUEST_INCREMENTS_FIELD`; see that field's own
|
||||
docstring for why, unlike concurrency, a hop that goes on to succeed
|
||||
(or is the chain's own final failure) is deliberately never refunded.
|
||||
"""
|
||||
logging_obj: Final = request_kwargs.get("litellm_logging_obj")
|
||||
model_call_details: Final = getattr(logging_obj, "model_call_details", None)
|
||||
|
|
@ -1541,6 +1608,9 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details)
|
||||
if release_keys:
|
||||
await self._release_keys(release_keys)
|
||||
stale_request_increments: Final = _pop_reservations(model_call_details, _PENDING_REQUEST_INCREMENTS_FIELD)
|
||||
if stale_request_increments:
|
||||
await self._release_keys(stale_request_increments)
|
||||
|
||||
async def _pop_pending_concurrency_keys(
|
||||
self, kwargs: Mapping[str, object]
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import (
|
|||
_fixed_length_identity,
|
||||
_inflight_key,
|
||||
_partition_key,
|
||||
_pending_reservations_cache_key,
|
||||
_PROXY_ModelBasedTagRateLimitsHook,
|
||||
)
|
||||
from litellm.types.router import RoutingGroup, TagRateLimitEntry, TagRateLimitScope
|
||||
|
|
@ -187,6 +188,26 @@ def test_fixed_length_identity_preserves_distinctness():
|
|||
assert _fixed_length_identity("user-a") == _fixed_length_identity("user-a")
|
||||
|
||||
|
||||
def test_pending_reservations_cache_key_bounds_call_id_regardless_of_input_size():
|
||||
"""
|
||||
veria-ai finding on PR #36541: litellm_call_id comes straight from the
|
||||
caller-controlled x-litellm-call-id header with no length bound, and was
|
||||
embedded directly in the pending-reservations mirror key -- a caller
|
||||
submitting long ids across many in-flight tagged requests could inflate
|
||||
Redis/in-memory key size disproportionately. Hashed via
|
||||
_fixed_length_identity, same as every other caller-controlled value this
|
||||
hook puts in a cache key.
|
||||
"""
|
||||
huge_call_id = "x" * 5_000_000
|
||||
key = _pending_reservations_cache_key(huge_call_id, "some-key-hash")
|
||||
assert len(key) < 200
|
||||
|
||||
|
||||
def test_pending_reservations_cache_key_preserves_distinctness():
|
||||
assert _pending_reservations_cache_key("call-a", "kh") != _pending_reservations_cache_key("call-b", "kh")
|
||||
assert _pending_reservations_cache_key("call-a", "kh") == _pending_reservations_cache_key("call-a", "kh")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_oversized_tag_value_does_not_inflate_the_bucket_key(time_controller):
|
||||
"""
|
||||
|
|
@ -2858,6 +2879,133 @@ async def test_next_hops_admission_releases_a_prior_hops_leaked_reservation(time
|
|||
assert result == healthy
|
||||
|
||||
|
||||
def _request_limit_router(limit: int) -> "litellm.Router":
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
_deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{
|
||||
"request_limits": {
|
||||
"limits": [{"name": "per_period", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300}]
|
||||
}
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_next_hops_admission_refunds_a_prior_failed_hops_request_increment(time_controller):
|
||||
"""
|
||||
Regression test for Cursor Bugbot's "fallback hops burn request budget"
|
||||
finding on PR #36541, live-confirmed against a real proxy: a "requests"
|
||||
limit is meant to cap logical client requests, not internal routing
|
||||
attempts, but without a refund a chain that fails once before succeeding
|
||||
burned 2 units of a 1-request-per-period budget for one logical call --
|
||||
live reproduction showed the retry's own admission rejected with
|
||||
current=1.0 limit=1.0 even though the client only made one call.
|
||||
|
||||
Concurrency's next-hop-releases-the-prior-hop's-stale-reservation pattern
|
||||
(see test_next_hops_admission_releases_a_prior_hops_leaked_reservation)
|
||||
generalizes cleanly here: since Router only re-enters admission for a hop
|
||||
that already failed, the prior hop's own "requests" increment must be
|
||||
refunded there too, before this hop's own check runs.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
router = _request_limit_router(limit=1)
|
||||
limiter.update_variables(llm_router=router)
|
||||
healthy = router.model_list
|
||||
request_kwargs, _kwargs = _call_context(["end_user_id:u1"])
|
||||
|
||||
# Hop 1 admits (the only unit) then fails -- no failure event follows,
|
||||
# mirroring the "already consumed litellm's one dedup-allowed failure
|
||||
# event" scenario the sibling concurrency test documents.
|
||||
result = await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
assert result == healthy
|
||||
|
||||
# Hop 2's own admission must refund hop 1's now-stale "requests"
|
||||
# increment before checking its own -- if it didn't, this raises
|
||||
# ProxyRateLimitError against a bucket a real client only asked to use
|
||||
# once.
|
||||
result = await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
assert result == healthy
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_next_hops_admission_refunds_a_request_increment_even_after_the_first_hops_own_failure_event_fires(
|
||||
time_controller,
|
||||
):
|
||||
"""
|
||||
Tighter regression than the test above: this reproduces the exact live
|
||||
failure this fix first shipped with. litellm's has_logged_async_failure
|
||||
dedup allows exactly the *first* failing hop's own async_log_failure_event
|
||||
through -- unlike a hop after that one, hop 1 here genuinely gets a real
|
||||
failure event, not silence. An earlier version of this fix popped
|
||||
_PENDING_REQUEST_INCREMENTS_FIELD in async_log_failure_event "for
|
||||
hygiene", discarding hop 1's entry before hop 2's own admission
|
||||
(_release_stale_hop_reservations) ever got a chance to refund it --
|
||||
silently and permanently stranding the charge, so hop 2 was rejected
|
||||
against a bucket a real client only asked to use once, live-confirmed
|
||||
against a real proxy. async_log_failure_event must leave this field
|
||||
completely untouched.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
router = _request_limit_router(limit=1)
|
||||
limiter.update_variables(llm_router=router)
|
||||
healthy = router.model_list
|
||||
request_kwargs, kwargs = _call_context(["end_user_id:u1"])
|
||||
|
||||
result = await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
assert result == healthy
|
||||
|
||||
# Hop 1's own, real failure event -- the one has_logged_async_failure
|
||||
# lets through.
|
||||
await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
|
||||
|
||||
# Hop 2's own admission must still refund hop 1's now-stale "requests"
|
||||
# increment before checking its own.
|
||||
result = await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
assert result == healthy
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_hops_own_request_increment_is_not_refunded(time_controller):
|
||||
"""
|
||||
The fix above must not swing the other way and refund every hop's
|
||||
"requests" increment unconditionally -- exactly one unit must survive
|
||||
per logical request, or the limit stops limiting anything. Simulates the
|
||||
full lifecycle (admission, then the success event a real request would
|
||||
fire) and confirms a second, unrelated logical request against the same
|
||||
tag is correctly rejected: the first request's own successful hop
|
||||
already spent the only unit for this period.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
router = _request_limit_router(limit=1)
|
||||
limiter.update_variables(llm_router=router)
|
||||
healthy = router.model_list
|
||||
request_kwargs, kwargs = _call_context(["end_user_id:u1"])
|
||||
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
|
||||
|
||||
fresh_request_kwargs, _fresh_kwargs = _call_context(["end_user_id:u1"])
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=fresh_request_kwargs
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_own_rejection_does_not_release_a_live_reservation(time_controller):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue