diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 66d0f0863f6..a7fd2a18fdc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index d48ae8fa643..92b905db49f 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -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 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 e48a1278cd7..97f1e065c68 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 @@ -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", + ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f6f0909e1a8..ad15941c992 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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): """