From edc548af97046f8604964549067cdf11982733dd Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 17:58:18 -0400 Subject: [PATCH] fix(proxy): require an unforgeable marker to trust litellm_metadata, and stop admission from releasing a live sibling's concurrency reservation Bugbot finding: resolve_authoritative_metadata_variable_name treated any non-empty litellm_metadata as authoritative, but a caller can populate it with unrelated keys on an ordinary route where "metadata" is the real field. Now requires the unconditionally-stamped, strip-protected "user_api_key_auth" marker instead of mere truthiness. Veria AI finding: Router.abatch_completion's comma-separated multi-model dispatch runs branches concurrently as separate asyncio Tasks that all share one litellm_logging_obj, so a still-live sibling branch's own concurrency reservation could sit in the same model_call_details a new hop's admission was cleaning up. _release_stale_hop_reservations now only reclaims entries queued by its own asyncio.Task, leaving a differently tasked (still-live) entry alone. --- .../hooks/model_based_tag_rate_limits_hook.py | 81 ++++++-- litellm/proxy/hooks/tag_rate_limits_shared.py | 19 +- .../test_model_based_tag_rate_limits_hook.py | 177 ++++++++++++++++-- 3 files changed, 239 insertions(+), 38 deletions(-) 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 45ae5943598..4c32f822b42 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -556,6 +556,15 @@ _INDEX_TTL_SECONDS: Final = 5.0 # server-side per logical request (and shared across that request's own # fallback hops, matching the original chain-wide release semantics), so it # can't be forged or guessed. +# +# "shared across that request's own fallback hops" is not the same as +# "scoped to one asyncio Task": `Router.abatch_completion`'s comma-separated +# multi-model dispatch runs several branches concurrently, each its own +# Task, but hands every branch the identical `litellm_logging_obj` -- so +# each entry also carries the Task that queued it (see +# `_queue_pending_reservations`), letting `_release_stale_hop_reservations` +# tell a genuinely stale same-task hop apart from a still-live sibling +# branch's own reservation. _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 @@ -893,6 +902,17 @@ def _queue_pending_reservations( logging object (defensive only; every real request has one): a queued concurrency reservation still self-heals via `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, just later. + + Each entry is stamped with the queueing coroutine's own `asyncio.Task`: + `Router.abatch_completion`'s comma-separated multi-model dispatch runs + several `acompletion` calls concurrently as *separate* tasks that all + share one `model_call_details` (the proxy attaches one `litellm_logging_obj` + to the request before the comma-split, and every branch inherits that + same reference), so this field is no longer scoped to one logical + request's own serial fallback chain the way its docstring assumes. + `_release_stale_hop_reservations` uses the stamp to tell "an earlier hop + of *this* chain, safe to reclaim" apart from "a concurrent sibling + branch's own still-live reservation," which must never be touched here. """ logging_obj: Final = request_kwargs.get("litellm_logging_obj") model_call_details: Final = getattr(logging_obj, "model_call_details", None) @@ -902,7 +922,10 @@ def _queue_pending_reservations( 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[field] = pending - pending.extend(reservations) # mutable-ok: see comment above + current_task: Final = asyncio.current_task() + pending.extend( + (key, partition_key, current_task) for key, partition_key in reservations + ) # mutable-ok: see comment above def _record_admission_time(request_kwargs: Mapping[str, object], now: float) -> None: @@ -1460,13 +1483,25 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] async def _release_stale_hop_reservations(self, request_kwargs: Mapping[str, object]) -> frozenset[str]: """ A concurrency reservation still queued when a *new* hop's admission - runs can only belong to an earlier hop of this same request that - already concluded and failed: Router awaits one hop's entire attempt - (call plus its own failure handling) before starting the next, and a - hop that instead succeeded ends the request there via - async_log_success_event, which already pops everything -- so - admission is never re-entered while an earlier hop's reservation is - still legitimately in flight. + runs, *within the same asyncio Task*, can only belong to an earlier + hop of this same request's own fallback chain that already concluded + and failed: Router awaits one hop's entire attempt (call plus its own + failure handling) before starting the next, and a hop that instead + succeeded ends the request there via async_log_success_event, which + already pops everything -- so admission is never re-entered, in that + same task, while an earlier hop's reservation is still legitimately + in flight. + + The task check matters because `model_call_details` is not always + scoped to one such chain: `Router.abatch_completion`'s comma-separated + multi-model dispatch runs several branches concurrently, each its own + Task, but every branch shares the identical `litellm_logging_obj` (see + `_queue_pending_reservations`'s own docstring) -- so a reservation + queued by a still-running sibling branch can be sitting here too, and + releasing it out from under that branch would let more calls through + a concurrency limit than it allows. Only entries this exact Task + queued are safe to treat as stale; anything else is left for its own + branch to release. LiteLLM only invokes a request's CustomLogger.async_log_failure_event once per request, for whichever hop fails first (its internal @@ -1506,16 +1541,17 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] model_call_details: Final = getattr(logging_obj, "model_call_details", None) if not isinstance(model_call_details, dict): return frozenset() - release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details) + release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_current_task=True) if release_keys: await self._release_keys(release_keys) pending_request_increments: Final = model_call_details.get(_PENDING_REQUEST_INCREMENTS_FIELD) if not isinstance(pending_request_increments, list): return frozenset() - return frozenset(key for key, _partition_key in pending_request_increments) + current_task: Final = asyncio.current_task() + return frozenset(key for key, _partition_key, task in pending_request_increments if task is current_task) async def _pop_pending_concurrency_keys( - self, kwargs: Mapping[str, object] + self, kwargs: Mapping[str, object], *, only_current_task: bool = False ) -> tuple[tuple[str, _PartitionKey], ...]: # Every caller of this method is itself a normal release path, so # also clear the async_post_call_failure_hook cache mirror for the @@ -1544,21 +1580,26 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] call_id, e, ) - # Snapshot then remove only those exact keys, never a blanket clear: - # a sibling hop sharing this same request's model_call_details can - # still be live and appending concurrently (see the field's own - # docstring), so wiping the whole list here would silently strand - # that hop's reservation instead of releasing it later. + # Snapshot then remove only those exact entries, never a blanket + # clear: a sibling branch sharing this same request's + # model_call_details can still be live and appending concurrently + # (see the field's own docstring), so wiping the whole list here + # would silently strand that branch's reservation instead of + # releasing it later. `only_current_task` additionally excludes any + # entry a *different*, still-running Task queued -- see + # `_release_stale_hop_reservations`'s own docstring for why that + # distinction, not just presence, decides what's actually stale. pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) if not isinstance(pending, list) or not pending: return () - keys: Final = tuple(pending) - for key in keys: + current_task: Final = asyncio.current_task() + snapshot: Final = tuple(entry for entry in pending if not only_current_task or entry[2] is current_task) + for entry in snapshot: try: - pending.remove(key) + pending.remove(entry) except ValueError: pass - return keys + return tuple(entry[:2] for entry in snapshot) async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: """ diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index 92c228ad819..dafc1b0a95d 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -232,14 +232,19 @@ def resolve_authoritative_metadata_variable_name( `metadata` dict for a standard (non LITELLM_METADATA_ROUTES) request -- and the key-presence check always picks `litellm_metadata` in both cases, silently reading no tags/identity at all and admitting the request against - every configured limit. Requiring the value to actually be a populated - dict, matching `_get_request_tags`'s own truthiness check in - litellm_logging.py, only ever prefers `litellm_metadata` when it is - genuinely the field the proxy wrote identity/tags into - (LITELLM_METADATA_ROUTES pre-seed it before admission runs, so it is - always a populated dict there, both at admission and at success time).""" + every configured limit. + + Merely requiring the value to be a non-empty dict is not enough either: a + caller can populate its own, unrelated keys on the non-authoritative + bucket (e.g. `{"litellm_metadata": {"x": 1}}` on an ordinary route), which + is non-empty but still not the field the proxy wrote identity into. + `add_litellm_data_to_request` unconditionally stamps `user_api_key_auth` + into whichever bucket the route actually resolved as authoritative, and + strips any `user_api_key_`-prefixed key a caller pre-populates on the + other bucket -- so requiring that marker's presence, not mere + truthiness, can't be forged onto the wrong side.""" litellm_metadata: Final = metadata_source.get("litellm_metadata") - if isinstance(litellm_metadata, Mapping) and litellm_metadata: + if isinstance(litellm_metadata, Mapping) and "user_api_key_auth" in litellm_metadata: return "litellm_metadata" return "metadata" 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 e9288566da0..60df49728f1 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 @@ -274,6 +274,39 @@ async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_presen ) +@pytest.mark.asyncio +async def test_filter_deployments_ignores_a_forged_populated_litellm_metadata_key(time_controller): + """ + Bugbot finding on the fix above: requiring litellm_metadata to be merely + non-empty is still forgeable -- a caller can populate it with its own, + unrelated keys on an ordinary route, which is non-empty but carries none + of the real identity the proxy wrote into "metadata". Only + add_litellm_data_to_request's own "user_api_key_auth" marker, stripped + from any bucket the caller doesn't own, proves a bucket is authoritative. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "litellm_metadata": {"x": 1}, + } + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # --------------------------------------------------------------------------- # TagRateLimitEntry -- limit validation # --------------------------------------------------------------------------- @@ -1438,6 +1471,54 @@ async def test_log_success_event_accounts_when_litellm_params_carries_a_null_lit ) +@pytest.mark.asyncio +async def test_log_success_event_ignores_a_forged_populated_litellm_metadata_key(time_controller): + """ + Same misresolution as the null-key case above, but with a populated + (not merely present) forged litellm_metadata -- a caller-supplied dict + with unrelated keys is still not the field the proxy wrote identity + into. Only the unconditionally-stamped "user_api_key_auth" marker, + which litellm_pre_call_utils.py strips from any bucket a caller doesn't + own, proves a bucket is authoritative. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "litellm_metadata": {"x": 1}, + "metadata": {"tags": ["end_user_id:u1"]}, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + @pytest.mark.asyncio async def test_log_success_event_accounts_the_key_backed_tag_not_a_caller_forged_one(time_controller): """ @@ -1523,7 +1604,7 @@ async def test_log_success_event_reads_nested_litellm_metadata_when_that_is_auth kwargs = { "litellm_params": { "metadata": {"tags": []}, - "litellm_metadata": {"tags": ["end_user_id:u1"]}, + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_auth": {}}, }, "standard_logging_object": { "model_group": "grp", @@ -1923,7 +2004,9 @@ async def test_log_success_event_accounts_against_the_team_id_admission_checked( # LITELLM_METADATA_ROUTES shape: litellm_metadata is the authoritative # field, and team-alias resolution requires the real team_id from it. - request_kwargs = {"litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}} + request_kwargs = { + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1", "user_api_key_auth": {}} + } result = await limiter.async_filter_deployments( model="team-alias-name", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs ) @@ -1933,7 +2016,9 @@ async def test_log_success_event_accounts_against_the_team_id_admission_checked( # real one in litellm_params.litellm_metadata, simulating litellm_logging.py # resolving a different field than the one admission used. kwargs = { - "litellm_params": {"litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}}, + "litellm_params": { + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1", "user_api_key_auth": {}} + }, "standard_logging_object": { "model_group": "team-alias-name", "model_id": "dep-1", @@ -2831,6 +2916,48 @@ async def test_next_hops_admission_releases_a_prior_hops_leaked_reservation(time assert result == healthy +@pytest.mark.asyncio +async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_controller): + """ + Veria AI finding: Router.abatch_completion's comma-separated multi-model + dispatch runs each model concurrently as its own asyncio.Task, but every + branch is handed the identical litellm_logging_obj (the proxy attaches + one to the request before the comma-split), so two genuinely concurrent + branches share one model_call_details. Before the fix, a second branch's + admission-time stale-reservation cleanup couldn't tell that apart from + an earlier, already-failed hop of its own retry chain, so it released + the first branch's still-live reservation and let both branches through + a concurrency limit of 1. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + first_admitted = asyncio.Event() + + async def _branch_one() -> None: + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + first_admitted.set() + # Held "in flight" while the second branch's admission runs, exactly + # like two concurrently in-flight provider calls. + await asyncio.sleep(0.05) + + async def _branch_two() -> None: + await first_admitted.wait() + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + results = await asyncio.gather( + asyncio.create_task(_branch_one()), asyncio.create_task(_branch_two()), return_exceptions=True + ) + rejections = [result for result in results if isinstance(result, ProxyRateLimitError)] + assert len(rejections) == 1 + + def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[ @@ -3985,13 +4112,16 @@ def test_concurrency_ttl_floor_does_not_shorten_a_longer_period_seconds(): @pytest.mark.asyncio async def test_release_in_a_forked_task_is_visible_to_the_parent_context(time_controller): limiter = _make_limiter(time_controller) - model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: ["key1"]} + # Entries are (key, partition_key, queueing_task) triples in production + # (see _queue_pending_reservations); the task is irrelevant to this + # specific release path (only_current_task defaults False here). + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} async def detached_release(): return await limiter._pop_pending_concurrency_keys(model_call_details) released = await asyncio.create_task(detached_release()) - assert released == ("key1",) + assert released == (("key1", None),) # The parent's own view of the same dict must see the release too. assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [] @@ -4000,31 +4130,56 @@ async def test_release_in_a_forked_task_is_visible_to_the_parent_context(time_co @pytest.mark.asyncio async def test_release_does_not_sweep_up_a_key_appended_after_its_snapshot(time_controller): limiter = _make_limiter(time_controller) - model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: ["key1"]} + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} async def detached_release_then_sibling_admits(): released = await limiter._pop_pending_concurrency_keys(model_call_details) # A sibling hop's admission, appending to the same shared dict, # interleaved right after this release's snapshot was taken. - model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD].append("key2") + model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD].append(("key2", None, None)) return released released = await asyncio.create_task(detached_release_then_sibling_admits()) - assert released == ("key1",) + assert released == (("key1", None),) # key2 must still be pending for its own hop's eventual release. - assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == ["key2"] + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("key2", None, None)] @pytest.mark.asyncio async def test_release_is_not_repeated_for_the_same_snapshot(time_controller): limiter = _make_limiter(time_controller) - model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: ["key1"]} + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} first = await limiter._pop_pending_concurrency_keys(model_call_details) second = await limiter._pop_pending_concurrency_keys(model_call_details) - assert first == ("key1",) + assert first == (("key1", None),) assert second == () +@pytest.mark.asyncio +async def test_release_only_current_task_leaves_a_concurrent_siblings_reservation_alone(time_controller): + """ + Veria AI finding: Router.abatch_completion's comma-separated multi-model + dispatch runs each model concurrently as its own asyncio.Task, but every + branch shares one litellm_logging_obj (the proxy attaches it to the + request before the comma-split), so a new hop's admission could see a + still-live sibling branch's own reservation sitting in the same + model_call_details and wrongly sweep it up as "stale". only_current_task + must leave a differently-tasked entry untouched. + """ + limiter = _make_limiter(time_controller) + + async def _reserve_as_a_separate_task() -> None: + pass # the task object itself is the fixture; body is irrelevant + + sibling_task = asyncio.create_task(_reserve_as_a_separate_task()) + await sibling_task + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("sibling-key", None, sibling_task)]} + + released = await limiter._pop_pending_concurrency_keys(model_call_details, only_current_task=True) + assert released == () + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_task)] + + # --------------------------------------------------------------------------- # refund-on-rollback across differently-hash-tagged keys (Redis Cluster safety) # ---------------------------------------------------------------------------