From 79b8d433dac19799e09198e911a0cf9e2b4a399d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 1 Sep 2026 10:47:31 -0400 Subject: [PATCH] fix(rate-limiting): check apply_to_models against each individual comma-separated model name apply_to_models was checked against the raw model field as a single string, so a request going through Router.abatch_completion (comma-separated model, e.g. "opus-chain,other") never matched a chain-scoped entry's apply_to_models list, since the joined string is never a member of a list of individual names. Confirmed live: a chain-scoped concurrency cap already at its limit for a single-model "opus-chain" request admitted a second "opus-chain,other" request without ever raising. _individual_model_names splits the caller's model field the same way route_llm_request.py does before it actually dispatches through Router.abatch_completion (call_type == acompletion, comma present). _classify's admission-time check and admitted_models' success-time accounting both now test candidate models individually via _entry_applies_any_candidate_model, a generalization of the existing admitted-models-only helper. _admitted_models_after now folds in every individual name from a batch admission, not the joined string, so the success-time accounting bug closes the same way the admission-time one does. --- .../hooks/global_tag_rate_limits_hook.py | 62 ++++++--- .../hooks/test_global_tag_rate_limits_hook.py | 121 ++++++++++++++++++ 2 files changed, 163 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index 67e8b679b6a..d095add3c77 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -144,18 +144,35 @@ else: Span: TypeAlias = object -def _entry_applies_any_admitted_model( - entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, admitted_models: frozenset[str] +def _entry_applies_any_candidate_model( + entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, candidate_models: frozenset[str] ) -> bool: """Same as `_entry_applies`, except an `apply_to_models`-scoped entry - counts as applying if ANY model an admission attempt for this call_id - saw was in scope -- not just whichever model the call ultimately served. - A `_pre_call_with_fallbacks` retry re-admits with a different model for - the same call_id, and an entry that matched an earlier attempt must - still get its success-time accounting.""" - if not admitted_models: + counts as applying if ANY of `candidate_models` is in scope -- not just a + single caller-visible name. Two call sites need this: at admission, a + comma-separated `model` (a non-racing or racing batch dispatch) names + several models at once, any of which should trip a chain-wide cap; at + success-event accounting, a `_pre_call_with_fallbacks` retry re-admits + with a different model for the same call_id, and an entry that matched + an earlier attempt must still get its accounting.""" + if not candidate_models: return _entry_applies(entry, tags, key_alias, None) - return any(_entry_applies(entry, tags, key_alias, model) for model in admitted_models) + return any(_entry_applies(entry, tags, key_alias, model) for model in candidate_models) + + +def _individual_model_names(model: str | None, call_type: str) -> tuple[str, ...]: + """`route_llm_request.py` splits a comma-separated `model` on this exact + condition before fanning out through `Router.abatch_completion`/ + `abatch_completion_fastest_response` -- an `apply_to_models`-scoped entry + must check each of those individual names, not the raw joined string + (which is never a member of any caller-configured `apply_to_models` + list), or a caller trivially bypasses a chain-wide cap by adding a + second model to the comma list.""" + if model is None: + return () + if call_type != "acompletion" or "," not in model: + return (model,) + return tuple(m.strip() for m in model.split(",")) def _non_racing_batch_width(data: Mapping[str, object], call_type: str) -> int: @@ -447,15 +464,17 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o @staticmethod def _admitted_models_after( - admitted_models: frozenset[str], model: str | None, renewal_allowed: bool + admitted_models: frozenset[str], candidate_models: frozenset[str], renewal_allowed: bool ) -> frozenset[str]: """Only called once this admission attempt has cleared every check - without raising -- a rejected attempt's model must never join + without raising -- a rejected attempt's models 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: - return admitted_models | frozenset((model,)) + this request under that model. `candidate_models` is every + individual name a comma-separated `model` names (see + `_individual_model_names`), not the raw joined string.""" + if renewal_allowed and candidate_models: + return admitted_models | candidate_models return admitted_models @staticmethod @@ -472,7 +491,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o key_alias: str | None, key_hash: str | None, now: float, - model: str | None, + candidate_models: frozenset[str], ) -> tuple[_ClassifiedGlobalCheck, ...]: classified: Final = [] # mutable-ok: sequential accumulator, immediately frozen into a tuple below for unit in _LIMIT_UNITS: @@ -483,7 +502,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o tag_value = _extract_identity(tags, entry.tag_id) if tag_value is None: continue - if not _entry_applies(entry, tags, key_alias, model): + if not _entry_applies_any_candidate_model(entry, tags, key_alias, candidate_models): continue effective_key_hash = key_hash if entry.scope_by_key_hash else None if unit == "concurrency": @@ -609,6 +628,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o 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 + candidate_models: Final = frozenset(_individual_model_names(model, call_type)) # First admission for this stash claims ownership; only a later one # with the same key_hash may renew its charges (see owner_key_hash). @@ -618,9 +638,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o now: Final = self._time_provider().timestamp() stash.admission_time = now - classified: Final = self._classify(config, tags, key_alias, key_hash, now, model) + classified: Final = self._classify(config, tags, key_alias, key_hash, now, candidate_models) if not classified: - stash.admitted_models = self._admitted_models_after(stash.admitted_models, model, renewal_allowed) + stash.admitted_models = self._admitted_models_after( + stash.admitted_models, candidate_models, renewal_allowed + ) return data read_only_checks: Final = tuple(c for c in classified if not c.is_atomic) @@ -700,7 +722,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 - stash.admitted_models = self._admitted_models_after(stash.admitted_models, model, renewal_allowed) + stash.admitted_models = self._admitted_models_after(stash.admitted_models, candidate_models, renewal_allowed) return data async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None: @@ -839,7 +861,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o tag_value = _extract_identity(tags, entry.tag_id) if tag_value is None: continue - if not _entry_applies_any_admitted_model(entry, tags, key_alias, admitted_models): + if not _entry_applies_any_candidate_model(entry, tags, key_alias, admitted_models): continue increment_value = increment_by_unit[unit] if increment_value == 0: diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 2867939a208..fffaf90baf7 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -610,6 +610,127 @@ async def test_apply_to_models_accounts_when_a_fallback_retry_re_admits_with_a_d ) +@pytest.mark.asyncio +async def test_apply_to_models_enforces_for_a_comma_separated_batch_model(time_controller, monkeypatch): + """Bugbot finding: a caller can dodge a chain-wide apply_to_models cap by + routing through Router.abatch_completion's comma-separated `model` -- + admission checks the individual names, not the raw joined string, which + would never match any configured apply_to_models entry.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"}, + call_type="acompletion", + ) + 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-2"), "model": "opus-chain,other-model"}, + call_type="acompletion", + ) + + +@pytest.mark.asyncio +async def test_apply_to_models_accounts_dollar_spend_from_a_batch_admission(time_controller, monkeypatch): + """Same gap as above, at success-time accounting: the batch admission's + own admitted_models must record each individual name, not the joined + string, or the spend never lands in the chain-scoped bucket at all.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "chain_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain,other-model"}, + call_type="acompletion", + ) + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": {"total_tokens": 0, "response_cost": 12.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + 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-2"), "model": "opus-chain"}, + call_type="acompletion", + ) + + +@pytest.mark.asyncio +async def test_apply_to_models_batch_split_requires_the_acompletion_call_type(time_controller, monkeypatch): + """route_llm_request.py only splits a comma-separated model for + call_type acompletion -- any other call type never reaches + Router.abatch_completion, so a comma there is just a literal (if odd) + model name and must not be split.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain,other-model"}, + call_type="embedding", + ) + assert result is not None + + @pytest.mark.asyncio async def test_a_rejected_admission_attempts_model_does_not_drive_later_accounting(time_controller, monkeypatch): """