scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix(rate-limiting): close three gaps in the global tag hook and fallback guard

Browse source 
- Wire order_tags_for_identity_resolution into global_tag_rate_limits_hook's
  admission and success-event tag resolution, matching the sibling hook.
  Without it a caller could put a forged tag ahead of the policy-backed
  inherited one and dodge or mis-bucket every global limit. Confirmed
  exploitable pre-fix, blocked post-fix, via a direct adversarial repro.
- Only record a model into the stash's admitted-models set once an
  admission attempt clears every check, not unconditionally at the top --
  a rejected attempt's model could otherwise still drive a later
  successful attempt's token/dollar accounting for an apply_to_models
  entry that never actually admitted the request.
- Check cross_model_scope on a fallback attempt's own rejection inside
  _pre_call_with_fallbacks's retry loop, not only on the original
  exception before the loop starts -- a chain-wide apply_to_models cap
  covering the first fallback too was previously bypassable by a second,
  uncovered fallback model.

Each fix has a regression test confirmed to fail on the pre-fix code and
pass on the fix.
This commit is contained in:
Deepanshu 2026-08-26 11:51:14 -04:00
parent 978fc694eb
commit 9a835e581f
4 changed files with 288 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/common_request_processing.py
View file

@ -2103,7 +2103,19 @@ class ProxyBaseLLMRequestProcessing:
route_type=route_type,
llm_router=llm_router,
)
except ProxyRateLimitError:
except ProxyRateLimitError as fallback_exc:
# A fallback attempt's own rejection can carry the
# identical cross_model_scope marker (this fallback
# model is itself covered by the same apply_to_models
# chain-wide cap) -- continuing to the next fallback
# would silently serve the request through a model
# outside that cap, defeating it just as much as not
# checking the original exception would.
if (
isinstance(fallback_exc.detail, Mapping)
and fallback_exc.detail.get("cross_model_scope") is True
):
raise
continue
except BaseException:
self.data["model"] = original_model

29
litellm/proxy/hooks/global_tag_rate_limits_hook.py
View file

@ -102,6 +102,9 @@ from litellm.proxy.hooks.tag_rate_limits_shared import (
from litellm.proxy.hooks.tag_rate_limits_shared import (
fixed_length_identity as _fixed_length_identity,
)
from litellm.proxy.hooks.tag_rate_limits_shared import (
order_tags_for_identity_resolution as _order_tags_for_identity_resolution,
)
from litellm.proxy.hooks.tag_rate_limits_shared import (
partition_key as _partition_key,
)
@ -391,6 +394,16 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
"global_tag_rate_limits_hook: failed to release concurrency slot %s: %s", key, e
)
@staticmethod
def _record_admitted_model(stash: _GlobalTagRateLimitStash, model: str | None, renewal_allowed: bool) -> None:
"""Only called once this admission attempt has cleared every check
without raising -- a rejected attempt's model must never join
admitted_models, or a later successful attempt's accounting could
wrongly credit an apply_to_models entry that never actually admitted
this request under that model."""
if renewal_allowed and model is not None:
stash.admitted_models = stash.admitted_models | frozenset((model,))
@staticmethod
def _ttl_for(unit: _LimitUnit, entry: TagRateLimitEntry) -> int:
if unit == "concurrency":
@ -525,7 +538,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
stash: Final = _claim_stash_for_data(data)
metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(data)
tags: Final = _get_tags_from_request_kwargs(data, metadata_variable_name=metadata_variable_name)
tags: Final = _order_tags_for_identity_resolution(
_get_tags_from_request_kwargs(data, metadata_variable_name=metadata_variable_name),
data,
metadata_variable_name,
)
key_alias: Final = user_api_key_dict.key_alias
key_hash: Final = user_api_key_dict.api_key
model: Final = data.get("model") if isinstance(data.get("model"), str) else None
@ -538,10 +555,9 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
now: Final = self._time_provider().timestamp()
stash.admission_time = now
if renewal_allowed and model is not None:
stash.admitted_models = stash.admitted_models | frozenset((model,))
classified: Final = self._classify(config, tags, key_alias, key_hash, now, model)
if not classified:
self._record_admitted_model(stash, model, renewal_allowed)
return data
read_only_checks: Final = tuple(c for c in classified if not c.is_atomic)
@ -613,6 +629,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
if request_keys:
stash.charged_request_keys.extend(request_keys) # mutable-ok: see field's own docstring
self._record_admitted_model(stash, model, renewal_allowed)
return data
async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None:
@ -662,7 +679,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name)
key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name)
tags: Final = _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)
tags: Final = _order_tags_for_identity_resolution(
_get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name),
kwargs,
metadata_variable_name,
)
if not tags:
return

186
tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py
View file

@ -590,6 +590,101 @@ async def test_apply_to_models_accounts_when_a_fallback_retry_re_admits_with_a_d
)
@pytest.mark.asyncio
async def test_a_rejected_admission_attempts_model_does_not_drive_later_accounting(time_controller, monkeypatch):
"""
A fallback retry's FIRST attempt can itself be rejected (by this same
entry, or a different hook) before it ever admits. That rejected
attempt's model must not join the stash's admitted-models history: a
later, successful attempt against a different (out-of-scope) model must
not have its accounting wrongly credited to an apply_to_models entry
that never actually admitted this request.
"""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [
{
"name": "conc-a",
"tag_id": "end_user_id",
"limit": 1,
"period_seconds": 60,
"apply_to_models": ["model-a"],
}
]
},
"dollar_limits": {
"limits": [
{
"name": "chain_spend",
"tag_id": "end_user_id",
"limit": 10.0,
"period_seconds": 86400,
"apply_to_models": ["model-a"],
}
]
},
},
)
hook = _make_hook(time_controller)
# Occupy model-a's only concurrency slot with an unrelated call.
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="occupier"), "model": "model-a"},
call_type="completion",
)
# call-1 attempt #1: model-a, rejected (slot taken) -- never truly admitted.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-a"},
call_type="completion",
)
# call-1 attempt #2 (fallback retry): model-b, not in apply_to_models=[model-a], admits.
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-b"},
call_type="completion",
)
kwargs = {
"litellm_call_id": "call-1",
"metadata": {"tags": ["end_user_id:u1"]},
"model": "model-b",
"standard_logging_object": {"total_tokens": 0, "response_cost": 50.0, "model": "model-b", "model_group": "model-b"},
}
await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# Release the occupier's concurrency slot so the final check below is
# gated only by chain_spend (dollars), not by conc-a still being full.
await hook.async_log_success_event(
kwargs={"litellm_call_id": "occupier", "metadata": {"tags": ["end_user_id:u1"]}},
response_obj=None,
start_time=0,
end_time=0,
)
await asyncio.sleep(0)
# chain_spend (apply_to_models=[model-a]) must still be empty: model-a's
# own admission attempt was rejected, never admitted, so a fresh
# model-a request is still allowed under the $100 limit.
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="call-2"), "model": "model-a"},
call_type="completion",
)
assert result is not None
# ---------------------------------------------------------------------------
# _pre_call_with_fallbacks reruns admission for the same logical request:
# a repeat call_id must renew, not double-charge -- veria-ai finding on
@ -1116,3 +1211,94 @@ async def test_config_reload_takes_effect_on_next_request(time_controller, monke
data=_data(["end_user_id:u1"], call_id="call-3"),
call_type="completion",
)
# ---------------------------------------------------------------------------
# Identity resolution: policy-backed tags must win over caller-supplied ones
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_a_caller_supplied_tag_cannot_shadow_the_policy_backed_identity_tag(time_controller, monkeypatch):
"""
_merge_tags (litellm_pre_call_utils.py) keeps caller-supplied tags first
in the merged tags list, appending key/team/project-contributed tags only
if not already present. Since extract_identity/entry_applies resolve a
tag_id by first-match-by-prefix, an authenticated caller could otherwise
submit e.g. company_id:attacker-chosen ahead of the key's real
company_id:real-company (surfaced via metadata.inherited_tags) and have
every company_id-scoped entry resolve to the forged value instead of the
real one -- letting the caller dodge the limit entirely by rotating
fabricated identities, or evade being charged against their own real
bucket. The hook must order tags so inherited_tags wins.
"""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"request_limits": {
"limits": [{"name": "per-company", "tag_id": "company_id", "limit": 1, "period_seconds": 86400}]
}
},
)
hook = _make_hook(time_controller)
key = _key()
poisoned_data = {
"litellm_call_id": "attack-1",
"metadata": {
"tags": ["company_id:attacker-chosen", "company_id:real-company"],
"inherited_tags": ["company_id:real-company"],
},
}
await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=poisoned_data, call_type="completion")
# The real company's own bucket must have been charged by the attack
# request, not a bucket keyed to the attacker's forged value -- so a
# second, genuine company_id:real-company request is now rejected.
victim_data = {
"litellm_call_id": "victim-1",
"metadata": {"tags": ["company_id:real-company"], "inherited_tags": ["company_id:real-company"]},
}
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=victim_data, call_type="completion")
@pytest.mark.asyncio
async def test_success_accounting_also_resolves_identity_from_the_policy_backed_tag(time_controller, monkeypatch):
"""Same forged-tag scenario as the admission-time test above, but for
async_log_success_event's own identity resolution (tokens/dollars)."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"dollar_limits": {
"limits": [{"name": "per-company-spend", "tag_id": "company_id", "limit": 10.0, "period_seconds": 86400}]
}
},
)
hook = _make_hook(time_controller)
kwargs = {
"litellm_call_id": "attack-1",
"metadata": {
"tags": ["company_id:attacker-chosen", "company_id:real-company"],
"inherited_tags": ["company_id:real-company"],
},
"standard_logging_object": {"total_tokens": 0, "response_cost": 20.0},
}
await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# The $20 spend must have landed against company_id:real-company, so a
# fresh request under the genuine identity is now over the $10 limit.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={
"litellm_call_id": "victim-1",
"metadata": {"tags": ["company_id:real-company"], "inherited_tags": ["company_id:real-company"]},
},
call_type="completion",
)

64
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -5865,6 +5865,70 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
assert call_count == 1
assert processor.data["model"] == primary_model
@pytest.mark.asyncio
async def test_cross_model_scoped_rejection_mid_chain_stops_further_fallback_attempts(self):
"""
Bugbot finding: the original exception is checked for
cross_model_scope before the fallback loop starts, but a LATER
fallback attempt's own rejection was never checked the same way --
the loop's `except ProxyRateLimitError: continue` swallowed it and
moved on to the next fallback model. If a chain-wide apply_to_models
cap covers both the primary model and the first fallback, and a
second fallback model isn't covered, this let the second fallback
silently serve the request the cap was meant to block. The original
(non-scoped) rejection enters the loop normally; the FIRST fallback's
own rejection carries cross_model_scope=True and must stop the loop
immediately, never reaching the second fallback.
"""
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
primary_model = "opus-chain"
processor = ProxyBaseLLMRequestProcessing(data={"model": primary_model})
attempted_models: list[str] = []
async def mock_pre_call_logic(**kwargs):
attempted_models.append(processor.data["model"])
if processor.data["model"] == primary_model:
# Original attempt: a plain, non-scoped rejection (e.g. a
# per-deployment limit), not the chain-wide cap itself.
raise ProxyRateLimitError(detail={"error": "tag_rate_limit_exceeded"}, headers={"retry-after": "30"})
if processor.data["model"] == "sonnet-chain":
# First fallback: rejected by the SAME chain-wide cap.
raise ProxyRateLimitError(
detail={"error": "tag_rate_limit_exceeded", "cross_model_scope": True},
headers={"retry-after": "30"},
)
raise AssertionError(f"must not attempt a second fallback model: {processor.data['model']}")
mock_router = MagicMock()
mock_router.fallbacks = [{"opus-chain": ["sonnet-chain", "haiku-chain"]}]
with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic):
with pytest.raises(ProxyRateLimitError) as exc_info:
await processor._pre_call_with_fallbacks(
request=MagicMock(),
general_settings={},
proxy_logging_obj=MagicMock(),
user_api_key_dict=MagicMock(router_settings=None),
version=None,
proxy_config=MagicMock(),
user_model=None,
user_temperature=None,
user_request_timeout=None,
user_max_tokens=None,
user_api_base=None,
model=primary_model,
route_type="acompletion",
llm_router=mock_router,
)
assert attempted_models == [primary_model, "sonnet-chain"]
assert exc_info.value.detail.get("cross_model_scope") is True
assert processor.data["model"] == primary_model
@pytest.mark.asyncio
async def test_real_parallel_request_limiter_model_tpm_limit_triggers_fallback(self):
"""