diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 3c1c6ef9d6e..b13e5719b73 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -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] diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index 9a429c26467..dffbe79ddc2 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -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): """