scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): require an unforgeable marker to trust litellm_metadata, and stop admission from releasing a live sibling's concurrency reservation

Browse source 
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.
This commit is contained in:
Deepanshu 2026-08-27 17:58:18 -04:00
parent bd673268d3
commit edc548af97
3 changed files with 239 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
litellm/proxy/hooks/model_based_tag_rate_limits_hook.py
View file

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

19
litellm/proxy/hooks/tag_rate_limits_shared.py
View file

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

177
tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py
View file

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