From d7629fc9c656f82680092b80bb73c6ee1e05fb41 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 10:32:54 -0400 Subject: [PATCH] feat(rate-limiting): add global tag rate limiting hook Adds global_tag_rate_limits_hook, a model-independent async_pre_call_hook that enforces tag rate limits across an entire fallback chain, plus a cross_model_scope fallback guard so a chain-wide rejection isn't silently retried against an unlisted fallback model. --- litellm/__init__.py | 3 + litellm/litellm_core_utils/litellm_logging.py | 26 + litellm/proxy/common_request_processing.py | 5 +- .../hooks/global_tag_rate_limits_hook.py | 789 +++++++++++++ .../hooks/test_global_tag_rate_limits_hook.py | 1005 +++++++++++++++++ .../proxy/test_common_request_processing.py | 59 + 6 files changed, 1886 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/hooks/global_tag_rate_limits_hook.py create mode 100644 tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 434e2fecff0..64006b37dc0 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -123,6 +123,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "dynamic_rate_limiter", "dynamic_rate_limiter_v3", "model_based_tag_rate_limits_hook", + "global_tag_rate_limits_hook", "langsmith", "prometheus", "otel", @@ -395,6 +396,8 @@ default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None model_based_tag_rate_limits_max_in_memory_cache_size: Optional[int] = None +global_tag_rate_limits: Optional["TagRateLimits"] = None +global_tag_rate_limits_max_in_memory_cache_size: Optional[int] = None model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e23ce66edd4..e8dbc25a58d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4556,6 +4556,23 @@ def _init_custom_logger_compatible_class( model_based_tag_rate_limits_hook_obj.update_variables(llm_router=llm_router) _in_memory_loggers.append(model_based_tag_rate_limits_hook_obj) return model_based_tag_rate_limits_hook_obj + elif logging_integration == "global_tag_rate_limits_hook": + from litellm.proxy.hooks.global_tag_rate_limits_hook import ( + _PROXY_GlobalTagRateLimitsHook, # pyright: ignore[reportPrivateUsage] # resolved by name like every other opt-in callback here + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_GlobalTagRateLimitsHook): + return callback + + if internal_usage_cache is None: + raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}") + + global_tag_rate_limits_hook_obj: Final = _PROXY_GlobalTagRateLimitsHook( + internal_usage_cache=internal_usage_cache + ) + _in_memory_loggers.append(global_tag_rate_limits_hook_obj) + return global_tag_rate_limits_hook_obj elif logging_integration == "langtrace": if "LANGTRACE_API_KEY" not in os.environ: raise ValueError("LANGTRACE_API_KEY not found in environment variables") @@ -5005,6 +5022,15 @@ def get_custom_logger_compatible_class( if isinstance(callback, _PROXY_ModelBasedTagRateLimitsHook): return callback + elif logging_integration == "global_tag_rate_limits_hook": + from litellm.proxy.hooks.global_tag_rate_limits_hook import ( + _PROXY_GlobalTagRateLimitsHook, # pyright: ignore[reportPrivateUsage] # resolved by name like every other opt-in callback here + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_GlobalTagRateLimitsHook): + return callback + elif logging_integration == "langtrace": from litellm.integrations.opentelemetry import OpenTelemetry diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 20059dcf3b4..66d0f0863f6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2061,7 +2061,10 @@ class ProxyBaseLLMRequestProcessing: ) except ProxyRateLimitError as original_exc: original_model: Final = self.data.get("model") - if not original_model or not llm_router or self.data.get("disable_fallbacks"): + cross_model_scope: Final = ( + isinstance(original_exc.detail, Mapping) and original_exc.detail.get("cross_model_scope") is True + ) + if not original_model or not llm_router or self.data.get("disable_fallbacks") or cross_model_scope: raise fallback_models: Final = self._resolve_fallback_models( diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py new file mode 100644 index 00000000000..ce8f59461d0 --- /dev/null +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -0,0 +1,789 @@ +""" +Tag-scoped token, request, dollar, and concurrency rate limits declared once, +globally, in `litellm_settings.global_tag_rate_limits` -- enforced once per +request in `async_pre_call_hook`, before Router does any routing, so a limit +applies regardless of which model or fallback chain the request ends up +hitting. + +This is the model-independent sibling of `model_based_tag_rate_limits_hook`, +which enforces the same `TagRateLimitEntry` shape but nested per-deployment +under `model_info.tag_rate_limits`, once per routing hop +(`async_filter_deployments`). A global entry has no deployment/routing-group +to reconcile -- there is exactly one config value, read once -- so this hook +reuses that sibling's free, already-hardened helper functions +(`entry_applies`, the Lua atomic check-and-increment scripts, cache +partitioning, bucket-key hashing primitives) from `tag_rate_limits_shared.py` +rather than duplicating them, but implements its own, much smaller +admission/accounting engine: no `_LimitsIndex`, no routing-group or +team-alias resolution, no per-deployment dedup signatures. + +Three independent entry-level knobs decide who a global entry applies to and +how its bucket is shared: + +- `apply_to_key_alias`: unset means every request, any key, any model. Set + to a list of virtual-key aliases, only those keys' requests count. +- `apply_to_models`: unset means every model. Set to a list of model names, + only requests whose caller-facing `model` field is in that list count -- + letting one entry rate-limit a whole fallback chain as a single unit by + naming every model in the chain. Each check is a fresh, independent + evaluation of `_entry_applies` against whatever `model` is current at that + moment, not a one-time decision that then sticks for the rest of the + request. Two concrete consequences follow from that: + (1) if the request's own model fails mid-flight and Router internally + retries a different model for the *same* admitted call, that retry is + never re-checked -- the original admission (against the originally + requested model) already stands, so an operator who needs the limit to + track whichever model actually ends up serving a request needs + `model_info.tag_rate_limits` instead; but + (2) if this hook's own admission *rejects* the request, + `common_request_processing.py` would otherwise catch that rejection and + retry the whole pre-call pipeline against + `litellm_settings.fallbacks`/`router_settings.fallbacks`, with + `data["model"]` mutated to the fallback target -- silently admitting the + request via a model outside `apply_to_models`, defeating the cap. A + rejection from an `apply_to_models`-scoped entry carries + `detail["cross_model_scope"] = True` for exactly this reason: + `_pre_call_with_fallbacks` checks that marker and re-raises immediately + instead of trying any fallback, so this bypass is closed regardless of + whether the fallback chain is also listed in `apply_to_models`. +- `scope_by_key_hash` (already exists on `TagRateLimitEntry`): whether the + keys an entry applies to share one bucket, or each gets its own. + +`async_pre_call_hook` runs before Router constructs `Logging`/`litellm_logging_obj` +for this request (see `common_request_processing.py`: `pre_call_hook` fires +well before `base_process_llm_request` builds the logging object), so unlike +`model_based_tag_rate_limits_hook` this hook cannot stash pending concurrency +reservations on `data["litellm_logging_obj"].model_call_details` -- that +object doesn't exist yet. Per-request state is instead kept on a +`ContextVar`-based stash, the same established pattern +`parallel_request_limiter_v3.py`'s v3 handler already uses for exactly this +problem, with one difference: the stash here is a dict keyed by +`litellm_call_id` rather than one shared mutable instance with an +overwritable "owner" field, so a nested LiteLLM call made inside the request +(e.g. a guardrail's own LLM judge call) -- which mints its own fresh call id +but inherits the same ContextVar-held ancestor context, not a separate one +-- gets its own isolated entry instead of overwriting the outer call's and +having its own success callback release the outer call's still-pending +reservations early. +""" + +import asyncio +from collections.abc import Callable, Mapping, Sequence +from contextvars import ContextVar +from dataclasses import dataclass, field +from datetime import datetime +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching model_based_tag_rate_limits_hook's identical import +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + ATOMIC_UNITS as _ATOMIC_UNITS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + BACKGROUND_TASKS as _BACKGROUND_TASKS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + LIMIT_UNITS as _LIMIT_UNITS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + TAG_RL_CHECK_AND_INCR_SCRIPT, + TAG_RL_DECR_FLOOR_ZERO_SCRIPT, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + UNIT_TO_GROUP_FIELD as _UNIT_TO_GROUP_FIELD, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + UNIT_TO_RATE_LIMIT_TYPE as _UNIT_TO_RATE_LIMIT_TYPE, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + LimitUnit as _LimitUnit, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + PartitionKey as _PartitionKey, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + PartitionOperations as _PartitionOperations, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + bucket_ttl_seconds as _bucket_ttl_seconds, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + entry_applies as _entry_applies, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + extract_identity as _extract_identity, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + extract_key_alias as _extract_key_alias, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + extract_key_hash as _extract_key_hash, +) +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 ( + partition_key as _partition_key, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + policy_fingerprint as _policy_fingerprint, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + resolve_success_event_metadata_variable_name as _resolve_success_event_metadata_variable_name, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.router_strategy.tag_based_routing import ( + _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching model_based_tag_rate_limits_hook's identical import +) +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import TagRateLimitEntry, TagRateLimits +from litellm.types.utils import StandardLoggingPayload + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span: TypeAlias = _Span +else: + Span: TypeAlias = object + + +def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: + """ + Global-hook equivalent of `model_based_tag_rate_limits_hook._hash_tag`, + without a `model_group`/deployment-scope/team-scope dimension -- a global + entry has none of those. Namespaced under `tag_rl:global:` so it can never + collide with that sibling hook's own `tag_rl:{model_group}:...` keys even + if an operator names a deployment "global": every key also differs by + `unit`/`name`/`tag_id`/`_policy_fingerprint`, and the two hooks' entries + are never meant to share a bucket in the first place. + """ + key_suffix: Final = f":key:{key_hash}" if key_hash is not None else "" + policy_suffix: Final = f":policy:{_policy_fingerprint(entry)}" + return f"tag_rl:global:{unit}:{entry.name}:{entry.tag_id}:{_fixed_length_identity(tag_value)}{key_suffix}{policy_suffix}" + + +def _bucket_key( + entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, bucket_id: int, key_hash: str | None +) -> str: + return f"{{{_hash_tag(entry, unit, tag_value, key_hash)}}}:{bucket_id}" + + +def _inflight_key(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: + return f"{{{_hash_tag(entry, unit, tag_value, key_hash)}}}:inflight" + + +@dataclass(frozen=True, slots=True) +class _ClassifiedGlobalCheck: + unit: _LimitUnit + entry: TagRateLimitEntry + tag_value: str + key: str + is_atomic: bool + + +@dataclass(frozen=True, slots=True) +class _CachePartition: + internal_usage_cache: InternalUsageCache + v3: _PROXY_MaxParallelRequestsHandler_v3 + + +@dataclass(slots=True) +class _GlobalTagRateLimitStash: + """Per-call bookkeeping `async_pre_call_hook` hands to that same call's + success/failure/disconnect callbacks -- see module docstring for why + this lives on a `ContextVar`, not `model_call_details`. + + Keyed by `litellm_call_id` in the dict below rather than one shared + mutable instance with an overwritable "owner" field: a nested LiteLLM + call made inside the request (an LLM-judge guardrail, a silent + experiment) that mints its own fresh call id runs inside the *same* + inherited context, not a separate one, so a single shared instance's + owner field would get reassigned to the nested call and its own + success callback would then release the outer call's still-pending + reservations early -- letting extra same-tag requests through while + the outer request is still genuinely in flight. Keying by call id + isolates each call's own reservations regardless of nesting. + """ + + admission_time: float | None = None + # The caller-facing `model` admission read from `data.get("model")`, so + # async_log_success_event's tokens/dollars accounting gates + # apply_to_models against the same, originally-requested model admission + # decided on -- not whatever model a later fallback actually served. + model: str | None = None + pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue + # "requests" keys already charged for this call_id -- veria-ai finding: + # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reruns the whole + # pre-call pipeline (this hook included) once per fallback model on ANY + # ProxyRateLimitError, not only one this hook itself raised, but reuses + # the same litellm_call_id (self.data is mutated in place, only `model` + # changes) across every attempt -- so this stash is the SAME object each + # time. A "requests" check matching an already-charged key here renews + # at zero net cost instead of charging a second unit for the same + # logical request; see async_pre_call_hook's own comment for how. + charged_request_keys: list[str] = field(default_factory=list) # mutable-ok: see comment above + # The server-authenticated key_hash (UserAPIKeyAuth.api_key) of whichever + # call first claimed this stash. litellm_call_id is caller-controlled via + # the x-litellm-call-id header (the exact forgery vector + # model_based_tag_rate_limits_hook's own pending-reservations mirror was + # hardened against earlier), so two unrelated requests sharing a + # caller-chosen id must not be allowed to "renew" each other's charge -- + # only a later admission carrying this same, authenticated key_hash may. + owner_key_hash: str | None = None + + +# Sentinel key for a call with no litellm_call_id at all (claim and lookup +# both fall back to this same key, so behavior for that degenerate case is +# unchanged: everything without a call id still shares one bucket). +_NO_CALL_ID: Final = "" + +_StashByCallId: TypeAlias = dict[ + str, _GlobalTagRateLimitStash +] # mutable-ok: per-call-id entries added over a request's lifetime, see class docstring + +_request_stash: Final[ContextVar[_StashByCallId | None]] = ContextVar( + "global_tag_rate_limits_request_stash", default=None +) + + +def _claim_stash_for_data(data: Mapping[str, object]) -> _GlobalTagRateLimitStash: + by_call_id: _StashByCallId | None = _request_stash.get() # rebind-ok: lazily initialized below if never set + if by_call_id is None: + by_call_id = {} # rebind-ok: see above # mutable-ok: see _StashByCallId + _request_stash.set(by_call_id) + owner_call_id: Final = data.get("litellm_call_id") + key: Final = owner_call_id if isinstance(owner_call_id, str) else _NO_CALL_ID + stash = by_call_id.get(key) # rebind-ok: reassigned just below when newly created + if stash is None: + stash = _GlobalTagRateLimitStash() # rebind-ok: see above + by_call_id[key] = stash # mutable-ok: see class docstring + return stash + + +def _stash_for_call(litellm_call_id: str | None) -> _GlobalTagRateLimitStash | None: + by_call_id: Final = _request_stash.get() + if by_call_id is None: + return None + key: Final = litellm_call_id if litellm_call_id is not None else _NO_CALL_ID + return by_call_id.get(key) + + +def _call_id_from_kwargs(kwargs: Mapping[str, object]) -> str | None: + call_id: Final = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def _resolve_max_in_memory_cache_size() -> int | None: + """Same shape as `model_based_tag_rate_limits_hook`'s own function, reading + this hook's own `litellm_settings` knob instead.""" + configured: Final = litellm.global_tag_rate_limits_max_in_memory_cache_size + if isinstance(configured, int) and not isinstance(configured, bool) and configured > 0: + return configured + if configured is not None: + verbose_proxy_logger.warning( + "global_tag_rate_limits_hook: global_tag_rate_limits_max_in_memory_cache_size=%r is not a positive " + "integer; falling back to the default in-memory cache size.", + configured, + ) + return None + + +class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # only referenced via the deferred import in litellm_logging.py's callback resolver; basedpyright doesn't trace that usage + CustomLogger +): + def __init__( + self, + internal_usage_cache: DualCache, + time_provider: Callable[[], datetime] | None = None, + ) -> None: + self._redis_cache: Final = internal_usage_cache.redis_cache + self._time_provider = time_provider or datetime.now + self._partitions: dict[_PartitionKey, _CachePartition] = {} # mutable-ok: lazily memoized; see _partition_for + self._partitions_lock = asyncio.Lock() + default_partition: Final = self._build_partition(_resolve_max_in_memory_cache_size()) + self._partitions[None] = default_partition + self.internal_usage_cache = default_partition.internal_usage_cache + self._lock = asyncio.Lock() + redis_cache: Final = self._redis_cache + self._check_and_incr_script = ( + redis_cache.async_register_script(TAG_RL_CHECK_AND_INCR_SCRIPT) if redis_cache is not None else None + ) + self._decr_floor_zero_script = ( + redis_cache.async_register_script(TAG_RL_DECR_FLOOR_ZERO_SCRIPT) if redis_cache is not None else None + ) + self._config_cache_key: object | None = None + self._config: TagRateLimits | None = None + + def _refresh_config(self) -> TagRateLimits | None: + """Re-validates `litellm.global_tag_rate_limits` whenever the object + identity changes (a config reload replaces it wholesale via + `setattr(litellm, key, value)`), so a hot-reloaded config takes effect + on the very next request with no staleness window and no TTL to tune.""" + raw: Final = getattr(litellm, "global_tag_rate_limits", None) + if raw is not self._config_cache_key: + self._config = TagRateLimits.model_validate(raw) if raw else None + self._config_cache_key = raw + return self._config + + def _build_partition(self, cache_size_override: int | None) -> _CachePartition: + dual_cache: Final = DualCache( + in_memory_cache=InMemoryCache(max_size_in_memory=cache_size_override), + redis_cache=self._redis_cache, + ) + cache: Final = InternalUsageCache(dual_cache=dual_cache) + return _CachePartition( + internal_usage_cache=cache, + v3=_PROXY_MaxParallelRequestsHandler_v3(cache, time_provider=self._time_provider), + ) + + async def _partition_for(self, partition_key: _PartitionKey) -> _CachePartition: + existing: Final = self._partitions.get(partition_key) + if existing is not None: + return existing + async with self._partitions_lock: + existing_after_lock: Final = self._partitions.get(partition_key) + if existing_after_lock is not None: + return existing_after_lock + cache_size_override: Final = partition_key[-1] if partition_key is not None else None + built: Final = self._build_partition(cache_size_override) + self._partitions[partition_key] = ( + built # mutable-ok: lazily memoized per distinct partition key, guarded by _partitions_lock above + ) + return built + + async def _check_and_increment_one( + self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int + ) -> tuple[bool, float]: + if self._check_and_incr_script is not None: + raw: Final = await self._check_and_incr_script(keys=(key,), args=(limit, increment, ttl)) + return bool(raw[0]), float(raw[1]) + async with self._lock: + current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) + current: Final = float(current_value) if current_value is not None else 0.0 + if current + increment > limit: + return False, current + new_value: Final = current + increment + await cache.async_set_cache(key=key, value=new_value, ttl=ttl, litellm_parent_otel_span=None) + return True, new_value + + async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None: + if self._decr_floor_zero_script is not None: + await self._decr_floor_zero_script(keys=(key,), args=(delta,)) + return + async with self._lock: + current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) + current: Final = float(current_value) if current_value is not None else 0.0 + await cache.async_set_cache(key=key, value=max(0.0, current + delta), litellm_parent_otel_span=None) + + async def _atomic_check_and_increment( + self, + checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], + ) -> tuple[int | None, tuple[float, ...]]: + """All-or-nothing atomic admission across `checks` -- see + `model_based_tag_rate_limits_hook._PROXY_ModelBasedTagRateLimitsHook._atomic_check_and_increment`'s + own docstring for the full rationale (refund-on-rollback, why a + raising key's own outcome is never refunded); identical logic, + duplicated rather than shared since it lives as instance methods + rather than free functions.""" + if not checks: + return None, () + admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection + for index, (cache, key, limit, increment, ttl) in enumerate(checks): + admitted = False + try: + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl) + finally: + if not admitted: + await self._refund_admitted(checks, up_to_index=index) + if admitted: + admitted_values.append(value) # mutable-ok: see accumulator comment above + continue + return index, (value,) + return None, tuple(admitted_values) + + async def _refund_admitted( + self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], up_to_index: int + ) -> None: + for refund_index in range(up_to_index): + refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index] + try: + await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment) + except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest + verbose_proxy_logger.warning( + "global_tag_rate_limits_hook: failed to refund %s on rollback: %s", refund_key, e + ) + + async def _release_keys(self, reservations: Sequence[tuple[str, _PartitionKey]]) -> None: + for key, partition_key in reservations: + try: + partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration + await self._decrement_floor_zero(partition.internal_usage_cache, key, -1.0) + except Exception as e: # noqa: BLE001 - releasing a slot must never raise into the caller's request path + verbose_proxy_logger.warning( + "global_tag_rate_limits_hook: failed to release concurrency slot %s: %s", key, e + ) + + @staticmethod + def _ttl_for(unit: _LimitUnit, entry: TagRateLimitEntry) -> int: + if unit == "concurrency": + requested_ttl: Final = entry.key_ttl_seconds if entry.key_ttl_seconds is not None else entry.period_seconds + return max(requested_ttl, _CONCURRENCY_MIN_SAFETY_TTL_SECONDS) + return _bucket_ttl_seconds(entry) + + def _classify( + self, + config: TagRateLimits, + tags: Sequence[str], + key_alias: str | None, + key_hash: str | None, + now: float, + model: str | None, + ) -> tuple[_ClassifiedGlobalCheck, ...]: + classified: Final = [] # mutable-ok: sequential accumulator, immediately frozen into a tuple below + for unit in _LIMIT_UNITS: + group = getattr(config, _UNIT_TO_GROUP_FIELD[unit]) + if group is None: + continue + for entry in group.limits: + tag_value = _extract_identity(tags, entry.tag_id) + if tag_value is None: + continue + if not _entry_applies(entry, tags, key_alias, model): + continue + effective_key_hash = key_hash if entry.scope_by_key_hash else None + if unit == "concurrency": + key = _inflight_key(entry, unit, tag_value, key_hash=effective_key_hash) + classified.append( + _ClassifiedGlobalCheck(unit, entry, tag_value, key, is_atomic=True) + ) # mutable-ok: see comment above + continue + bucket_id = int(now) // entry.period_seconds + key = _bucket_key(entry, unit, tag_value, bucket_id, key_hash=effective_key_hash) + classified.append( # mutable-ok: see comment above + _ClassifiedGlobalCheck(unit, entry, tag_value, key, is_atomic=unit in _ATOMIC_UNITS) + ) + return tuple(classified) + + async def _read_only_values( + self, read_only_checks: Sequence[_ClassifiedGlobalCheck], parent_otel_span: Span | None + ) -> tuple[float | None, ...]: + if not read_only_checks: + return () + indices_by_partition: Final[dict[_PartitionKey, list[int]]] = {} # mutable-ok: grouped, reassembled below + for index, check in enumerate(read_only_checks): + partition_key = _partition_key(check.entry) + indices = indices_by_partition.setdefault(partition_key, []) # mutable-ok: see above + indices.append(index) # mutable-ok: see comment above + values_by_index: Final[dict[int, float | None]] = {} # mutable-ok: see comment above + for partition_key, indices in indices_by_partition.items(): + partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration + keys = [read_only_checks[i].key for i in indices] # mutable-ok: async_batch_get_cache needs a real list + redis_cache = partition.internal_usage_cache.dual_cache.redis_cache + if redis_cache is not None: + redis_values: Mapping[str, object] = await redis_cache.async_batch_get_cache( + key_list=keys, parent_otel_span=parent_otel_span + ) + resolved = [redis_values.get(key) for key in keys] # mutable-ok: needs a real list + else: + current_values = await partition.internal_usage_cache.async_batch_get_cache( + keys=keys, parent_otel_span=parent_otel_span, local_only=True + ) + missing = [None] * len(keys) # mutable-ok: async_batch_get_cache requires a real list; see above + resolved = current_values if current_values is not None else missing + for i, value in zip(indices, resolved): + values_by_index[i] = value # mutable-ok: see comment above + return tuple(values_by_index[i] for i in range(len(read_only_checks))) + + def _raise_if_over_limit( + self, + read_only_checks: Sequence[_ClassifiedGlobalCheck], + current_values: Sequence[float | None], + model: str | None, + ) -> None: + for check, current_value in zip(read_only_checks, current_values): + current = float(current_value) if current_value is not None else 0.0 + if current < check.entry.limit: + continue + self._raise_over_limit(check.unit, check.entry, check.tag_value, model, current=current) + + def _raise_over_limit( + self, unit: _LimitUnit, entry: TagRateLimitEntry, tag_value: str, model: str | None, current: float + ) -> None: + verbose_proxy_logger.debug( + "global_tag_rate_limits_hook: OVER_LIMIT unit=%s name=%s tag_id=%s tag_value=%s current=%s limit=%s", + unit, + entry.name, + entry.tag_id, + tag_value, + current, + entry.limit, + ) + raise ProxyRateLimitError( + detail={ # mutable-ok: async_log_failure_event and generic proxy exception rendering branch on isinstance(exc.detail, dict) + "error": "tag_rate_limit_exceeded", + "type": unit, + "tag_id": entry.tag_id, + "tag_value": tag_value, + "limit_name": entry.name, + "limit": entry.limit, + "period_seconds": entry.period_seconds, + # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reads this: + # an apply_to_models entry caps an entire named chain as one unit, so + # retrying against a fallback model outside that list would silently + # defeat the very policy that just rejected this request. + **({"cross_model_scope": True} if entry.apply_to_models is not None else {}), # mutable-ok: see above + }, + headers={"retry-after": str(entry.period_seconds)}, # mutable-ok: same as detail + rate_limit_type=_UNIT_TO_RATE_LIMIT_TYPE[unit], + model=model, + llm_provider="litellm_proxy", + ) + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, # mutable-ok: must match CustomLogger.async_pre_call_hook's own base signature exactly + call_type: str, + ) -> dict: # mutable-ok: must match CustomLogger.async_pre_call_hook's own base signature exactly + config: Final = self._refresh_config() + if config is None: + return data + + # async_pre_call_hook fires once per request in the common case, but + # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks can re-run + # this same pipeline once per fallback model on any ProxyRateLimitError + # (not only one this hook raised) -- see charged_request_keys' own + # docstring for how a repeat run for the same call_id renews rather + # than re-charges both "requests" and "concurrency" checks below. + 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) + 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 + + # Only a repeat admission carrying the SAME authenticated key_hash as + # whichever call first claimed this stash may renew its charges -- + # see owner_key_hash's own docstring for why a bare call_id match is + # not enough. First admission for this stash claims ownership here. + if stash.owner_key_hash is None: + stash.owner_key_hash = key_hash + renewal_allowed: Final = stash.owner_key_hash == key_hash + + now: Final = self._time_provider().timestamp() + stash.admission_time = now + stash.model = model + classified: Final = self._classify(config, tags, key_alias, key_hash, now, model) + if not classified: + return data + + read_only_checks: Final = tuple(c for c in classified if not c.is_atomic) + atomic_checks: Final = tuple(c for c in classified if c.is_atomic) + + current_values: Final = await self._read_only_values(read_only_checks, parent_otel_span=None) + self._raise_if_over_limit(read_only_checks, current_values, model) + + if atomic_checks: + atomic_partitions_list: Final = [] # mutable-ok: sequential async lookups, one per atomic_checks entry + for check in atomic_checks: + atomic_partitions_list.append( + await self._partition_for(_partition_key(check.entry)) + ) # mutable-ok: see comment above + atomic_partitions: Final = tuple(atomic_partitions_list) + already_reserved_concurrency_keys: Final = frozenset( + key for key, _partition_key in stash.pending_concurrency_keys + ) + failing_index, values = await self._atomic_check_and_increment( + tuple( + ( + partition.internal_usage_cache, + check.key, + check.entry.limit, + # A key already charged/reserved for this call_id (an + # earlier _pre_call_with_fallbacks attempt for the + # same logical request) renews at zero net cost + # instead of charging or reserving a second unit -- + # folded into this same all-or-nothing batch so a + # rollback here (some other check in the batch + # rejecting) refunds that zero-cost renewal as a + # genuine no-op, same reasoning as + # model_based_tag_rate_limits_hook's identical fix + # for its own per-hop retries. + 0.0 + if renewal_allowed + and ( + (check.unit == "requests" and check.key in stash.charged_request_keys) + or (check.unit == "concurrency" and check.key in already_reserved_concurrency_keys) + ) + else 1.0, + self._ttl_for(check.unit, check.entry), + ) + for partition, check in zip(atomic_partitions, atomic_checks) + ) + ) + if failing_index is not None: + failing_check: Final = atomic_checks[failing_index] + self._raise_over_limit( + failing_check.unit, failing_check.entry, failing_check.tag_value, model, current=values[0] + ) + + # Only genuinely new reservations, never a key already in + # already_reserved_concurrency_keys: that key's own check just + # renewed at zero net cost above, so re-adding it here would + # make release (which decrements once per queued entry) decrement + # twice for a counter that was only ever incremented once. + concurrency_reservations: Final = tuple( + (check.key, _partition_key(check.entry)) + for check in atomic_checks + if check.unit == "concurrency" and check.key not in already_reserved_concurrency_keys + ) + if concurrency_reservations: + stash.pending_concurrency_keys.extend(concurrency_reservations) # mutable-ok: see field's own docstring + + # Only recorded when renewal_allowed: an admission that didn't + # own this stash (a call_id collision from a different key_hash) + # must not contaminate the rightful owner's own renewal + # tracking, or a later, genuine fallback retry from the owner + # could wrongly treat the impostor's charge as its own and + # renew for free. + request_keys: Final = ( + tuple( + check.key + for check in atomic_checks + if check.unit == "requests" and check.key not in stash.charged_request_keys + ) + if renewal_allowed + else () + ) + if request_keys: + stash.charged_request_keys.extend(request_keys) # mutable-ok: see field's own docstring + + return data + + async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + stash: Final = _stash_for_call(_call_id_from_kwargs(request_data)) + if stash is None or not stash.pending_concurrency_keys: + return + release_keys: Final = tuple(stash.pending_concurrency_keys) + stash.pending_concurrency_keys.clear() + await self._release_keys(release_keys) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + # No special-case skip for this hook's own tag_rate_limit_exceeded + # rejection: that rejection never reaches the point where a + # concurrency reservation is queued (see async_pre_call_hook), so + # stash.pending_concurrency_keys is already empty in that case and + # the check below naturally no-ops. Skipping release based on the + # exception's error marker alone would be wrong here, since + # model_based_tag_rate_limits_hook raises the identical marker -- + # that rejection can land after this hook already reserved a slot + # for this same request, and that slot must still be released. + stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) + if stash is None or not stash.pending_concurrency_keys: + return + release_keys: Final = tuple(stash.pending_concurrency_keys) + stash.pending_concurrency_keys.clear() + await self._release_keys(release_keys) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) + if stash is not None and stash.pending_concurrency_keys: + release_keys: Final = tuple(stash.pending_concurrency_keys) + stash.pending_concurrency_keys.clear() + release_task: Final = asyncio.create_task(self._release_keys(release_keys)) + _BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + release_task.add_done_callback(_BACKGROUND_TASKS.discard) + + config: Final = self._refresh_config() + if config is None: + return + + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") + if standard_logging_object is None: + return + + # kwargs here is Logging.model_call_details, not the router's flat + # request kwargs admission sees: metadata/litellm_metadata are never + # top-level here, only nested under kwargs["litellm_params"] (see + # Logging.update_environment_variables). + litellm_params_for_metadata: Final = kwargs.get("litellm_params") or kwargs + metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata) + 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) + if not tags: + return + + now: Final = ( + stash.admission_time + if stash is not None and stash.admission_time is not None + else self._time_provider().timestamp() + ) + model: Final = stash.model if stash is not None else None + increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType( + { + "tokens": float(standard_logging_object.get("total_tokens") or 0), + "dollars": float(standard_logging_object.get("response_cost") or 0), + } + ) + + operation_by_entry: Final = [] # mutable-ok: sequential accumulator over config groups, immediately used below + for unit in ("tokens", "dollars"): + group = getattr(config, _UNIT_TO_GROUP_FIELD[unit]) + if group is None: + continue + for entry in group.limits: + tag_value = _extract_identity(tags, entry.tag_id) + if tag_value is None: + continue + if not _entry_applies(entry, tags, key_alias, model): + continue + increment_value = increment_by_unit[unit] + if increment_value == 0: + continue + bucket_id = int(now) // entry.period_seconds + key_hash_for_entry = key_hash if entry.scope_by_key_hash else None + key = _bucket_key(entry, unit, tag_value, bucket_id, key_hash=key_hash_for_entry) + operation_by_entry.append( # mutable-ok: see comment above + ( + entry, + RedisPipelineIncrementOperation( + key=key, increment_value=increment_value, ttl=_bucket_ttl_seconds(entry) + ), + ) + ) + + if not operation_by_entry: + return + + operations_by_partition: Final[_PartitionOperations] = {} # mutable-ok: grouped by cache partition below + for entry, operation in operation_by_entry: + partition_key = _partition_key(entry) + operations = operations_by_partition.setdefault(partition_key, []) # mutable-ok: see above + operations.append(operation) # mutable-ok: see comment above + + for partition_key, group_operations in operations_by_partition.items(): + partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration + accounting_task = asyncio.create_task( # not Final: rebound each loop iteration + partition.v3.async_increment_tokens_with_ttl_preservation( + pipeline_operations=tuple(group_operations), parent_otel_span=None + ) + ) + _BACKGROUND_TASKS.add(accounting_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + accounting_task.add_done_callback(_BACKGROUND_TASKS.discard) 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 new file mode 100644 index 00000000000..428037986e5 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -0,0 +1,1005 @@ +""" +Unit tests for the global-scope, model-independent tag rate limiter. +""" + +import asyncio +from datetime import datetime, timedelta + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.global_tag_rate_limits_hook import ( + _PROXY_GlobalTagRateLimitsHook, +) + + +class TimeController: + def __init__(self): + self._current = datetime(2026, 1, 1, 0, 0, 0) + + def now(self) -> datetime: + return self._current + + def advance(self, seconds: float) -> None: + self._current += timedelta(seconds=seconds) + + +@pytest.fixture +def time_controller(): + return TimeController() + + +def _make_hook(time_controller: TimeController) -> _PROXY_GlobalTagRateLimitsHook: + return _PROXY_GlobalTagRateLimitsHook( + internal_usage_cache=DualCache(), + time_provider=time_controller.now, + ) + + +def _key(alias: str | None = None, api_key: str = "hash") -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key=api_key, key_alias=alias) + + +def _data(tags: list[str], call_id: str = "call-1") -> dict: + return {"metadata": {"tags": tags}, "litellm_call_id": call_id} + + +# --------------------------------------------------------------------------- +# No-op when unconfigured +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_no_op_when_no_config_set(time_controller, monkeypatch): + monkeypatch.setattr(litellm, "global_tag_rate_limits", None) + 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_type="completion" + ) + assert result == _data(["end_user_id:u1"]) + + +@pytest.mark.asyncio +async def test_malformed_config_raises_at_first_use(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, "global_tag_rate_limits", {"dollar_limits": {"limits": [{"name": "bad", "limit": "not-a-number"}]}} + ) + hook = _make_hook(time_controller) + with pytest.raises(ValidationError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + + +# --------------------------------------------------------------------------- +# Global scope: applies to every key by default +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_request_limit_shared_across_keys_by_default(time_controller, monkeypatch): + """No apply_to_key_alias -> the entry is one shared bucket regardless of + which key made the request.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a"), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + # A different key, identical tag value, and a distinct call_id (a + # genuinely separate logical request, not a fallback retry of the same + # one) -- must be rejected too, proving the bucket is genuinely shared, + # not per-key by default. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-b"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_request_limit_is_independent_of_model(time_controller, monkeypatch): + """The hook never reads `data["model"]` for identity -- two different + "models" (irrelevant to this hook) must still share the same bucket.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + data_model_a = {**_data(["end_user_id:u1"]), "model": "gpt-4o"} + # A distinct call_id: this is a separate logical request, not the same + # one retrying against a different model via _pre_call_with_fallbacks. + data_model_b = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "claude-3"} + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=data_model_a, call_type="completion" + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=data_model_b, call_type="completion" + ) + + +# --------------------------------------------------------------------------- +# apply_to_key_alias -- narrows which keys an entry applies to +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_ignores_non_matching_keys(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + for _ in range(3): + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="other-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_enforces_for_the_listed_key(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + # A distinct call_id: a second, separate request from the same key, not + # a fallback retry of the first. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_composes_with_scope_by_key_hash(time_controller, monkeypatch): + """Both listed keys are subject to the entry, but scope_by_key_hash + splits their buckets: exhausting one must not affect the other.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_key_alias": ["key-a", "key-b"], + "scope_by_key_hash": True, + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a", api_key="hashA"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + # A distinct call_id: a second, separate request from the same key, not + # a fallback retry of the first. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a", api_key="hashA"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + # key-b is unaffected by key-a's exhausted bucket. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-b", api_key="hashB"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + assert result is not None + + +# --------------------------------------------------------------------------- +# apply_to_models -- narrows which requested model an entry applies to, +# letting one entry rate-limit a whole fallback chain as a single unit +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_to_models_ignores_non_matching_model(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + for i in range(3): + data = {**_data(["end_user_id:u1"], call_id=f"call-{i}"), "model": "sonnet-chain"} + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion" + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_apply_to_models_enforces_for_the_listed_model(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data1 = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data1, call_type="completion") + data2 = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"} + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data2, call_type="completion") + + +@pytest.mark.asyncio +async def test_apply_to_models_shares_one_bucket_across_every_listed_model(time_controller, monkeypatch): + """The core "rate limit the whole chain" use case: a single limit shared + across every model named in apply_to_models, not one bucket per model.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain", "sonnet-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data1 = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data1, call_type="completion") + + data2 = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "sonnet-chain"} + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data2, call_type="completion") + + +@pytest.mark.asyncio +async def test_apply_to_models_composes_with_apply_to_key_alias(time_controller, monkeypatch): + """Both gates must pass -- the listed key requesting a non-listed model + is unaffected, and only the listed key requesting the listed model is + actually enforced.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + # premium-key requesting a non-listed model: apply_to_models alone must + # still exclude it, even though apply_to_key_alias matches. + for i in range(3): + data = {**_data(["end_user_id:u1"], call_id=f"wrong-model-{i}"), "model": "sonnet-chain"} + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), cache=DualCache(), data=data, call_type="completion" + ) + assert result is not None + + # A non-listed key requesting the listed model: apply_to_key_alias alone + # must still exclude it, even though apply_to_models matches. + for i in range(3): + data = {**_data(["end_user_id:u1"], call_id=f"wrong-key-{i}"), "model": "opus-chain"} + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="other-key"), cache=DualCache(), data=data, call_type="completion" + ) + assert result is not None + + # Both gates match: enforced. + data1 = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), cache=DualCache(), data=data1, call_type="completion" + ) + data2 = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"} + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), cache=DualCache(), data=data2, call_type="completion" + ) + + +@pytest.mark.asyncio +async def test_dollar_limit_respects_apply_to_models_at_accounting_time(time_controller, monkeypatch): + """The entry only applies to opus-chain; a non-listed model's spend must + not be charged against this bucket at all -- proves apply_to_models + gates async_log_success_event's tokens/dollars accounting, not just + admission.""" + 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) + + data = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "sonnet-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": {"total_tokens": 0, "response_cost": 999.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # opus-chain was never charged -- still fully under its own 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": "opus-chain"}, + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_apply_to_models_fallback_does_not_re_narrow_accounting_to_the_serving_model(time_controller, monkeypatch): + """ + Documented limitation, not a bug: apply_to_models is evaluated exactly + once, at admission, against the caller-requested model -- it is never + re-evaluated against whichever model a later fallback actually serves. + This request names "opus-chain" at admission (the entry applies), but its + response accounting reports "sonnet-chain" as the model that actually + served it, simulating Router falling back after opus-chain failed. The + spend must still land in the opus-chain-scoped bucket: the check already + ran and decided at admission, and is not re-run for the fallback target. + """ + 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) + + data = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + + # The response accounting reports the fallback target as the model that + # actually served the request -- not the "opus-chain" admission decided. + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "model": "sonnet-chain", + "standard_logging_object": { + "total_tokens": 0, + "response_cost": 12.0, + "model": "sonnet-chain", + "model_group": "sonnet-chain", + }, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # The spend landed in the opus-chain-scoped bucket regardless -- a fresh + # opus-chain request is now over the limit. + 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="completion", + ) + + +# --------------------------------------------------------------------------- +# _pre_call_with_fallbacks reruns admission for the same logical request: +# a repeat call_id must renew, not double-charge -- veria-ai finding on +# PR #36541 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_repeat_admission_for_the_same_call_id_and_key_renews_instead_of_double_charging( + time_controller, monkeypatch +): + """ + ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reruns the whole + pre-call pipeline (this hook included) once per fallback model on ANY + ProxyRateLimitError, not only one this hook itself raised, but keeps the + same litellm_call_id across every attempt (self.data is mutated in + place; only "model" changes). Without this fix, an unrelated rejection + (a different rate limiter, a budget cap) triggering N fallback attempts + would charge this hook's own "requests" cap N times for one logical + client call. A limit of 1 makes a double-charge directly observable: if + the second admission (same call_id, same key) charged again instead of + renewing, this would raise. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + key = _key(alias="key-a") + + await hook.async_pre_call_hook( + user_api_key_dict=key, cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + # Same call_id, same key, different model -- exactly what + # _pre_call_with_fallbacks produces for a fallback attempt of the same + # logical request. + result = await hook.async_pre_call_hook( + user_api_key_dict=key, + cache=DualCache(), + data={**_data(["end_user_id:u1"]), "model": "fallback-model"}, + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_a_forged_shared_call_id_from_a_different_key_does_not_get_a_free_renewal(time_controller, monkeypatch): + """ + Security regression: litellm_call_id is caller-controlled via the + x-litellm-call-id header (the same forgery vector + model_based_tag_rate_limits_hook's own pending-reservations mirror was + hardened against earlier in this PR). Two unrelated requests choosing + the identical call_id must not be able to renew each other's charge -- + only a second admission carrying the SAME authenticated key_hash as + whichever request first claimed that call_id may. A limit of 1 makes + this observable: if the second, different-key admission wrongly + renewed, it would succeed instead of raising. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a", api_key="hashA"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="forged-call-id"), + call_type="completion", + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-b", api_key="hashB"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="forged-call-id"), + call_type="completion", + ) + + +# --------------------------------------------------------------------------- +# Concurrency: reservation at admission, release on success/failure/disconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_concurrency_limit_rejects_second_admission_until_release(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + 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"), + call_type="completion", + ) + 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"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_on_success(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + async def one_request(call_id: str) -> None: + data = _data(["end_user_id:u1"], call_id=call_id) + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = {"litellm_call_id": call_id, "metadata": {"tags": ["end_user_id:u1"]}} + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + await one_request("call-1") + await asyncio.sleep(0) # let the fire-and-forget release task run + + # The slot was released, so a fresh request must be admitted again. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_nested_call_success_does_not_release_the_outer_calls_reservation(time_controller, monkeypatch): + """ + A nested LiteLLM call made inside the request (an LLM-judge guardrail, a + silent experiment) mints its own fresh litellm_call_id but inherits the + same ContextVar-held stash as the outer call, since it runs in the same + task rather than a separate one. The outer call's own concurrency + reservation must survive the nested call's admission and success + callback: it belongs to a different call id and must not be touched by + it. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + 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-outer"), + call_type="completion", + ) + + # Nested call, same task, different tag and a fresh call id -- admits + # and completes entirely before the outer call's own success/failure + # callback ever fires. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u2"], call_id="call-nested"), + call_type="completion", + ) + await hook.async_log_success_event( + kwargs={"litellm_call_id": "call-nested", "metadata": {"tags": ["end_user_id:u2"]}}, + response_obj=None, + start_time=0, + end_time=0, + ) + await asyncio.sleep(0) + + # The outer call is still genuinely in flight -- its own reservation + # must still be held, so a second end_user_id:u1 request is rejected. + 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-outer-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_when_a_different_hook_rejects_the_request(time_controller, monkeypatch): + """ + model_based_tag_rate_limits_hook raises the identical ProxyRateLimitError + shape (detail["error"] == "tag_rate_limit_exceeded") this hook's own + admission does, since both hooks share the same rejection marker. + async_log_failure_event fires on every registered CustomLogger regardless + of which one raised, so this hook must still release its own successfully + reserved concurrency slot when the *other* hook is what rejected the + request -- skipping release just because the marker matches would leak + this hook's own slot until the safety TTL, even though nothing about this + hook's own admission failed. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + {"concurrency_limits": {"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]}}, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + + other_hooks_rejection = ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "type": "requests", "tag_id": "end_user_id"}, + headers={"retry-after": "60"}, + rate_limit_type=None, + model="gpt-4o", + llm_provider="litellm_proxy", + ) + kwargs = {"litellm_call_id": "call-1", "exception": other_hooks_rejection} + await hook.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # The slot was released despite the shared rejection marker, so a fresh + # request must be admitted again. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_on_disconnect(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) + + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_concurrent_requests_do_not_share_each_others_reservation_state(time_controller, monkeypatch): + """Two logically distinct requests running as separate asyncio Tasks must + not see each other's pending-concurrency stash, even though both share + this hook instance -- the whole point of the ContextVar-based stash.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 5, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + async def one_request(call_id: str) -> int: + data = _data(["end_user_id:u1"], call_id=call_id) + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = {"litellm_call_id": call_id, "metadata": {"tags": ["end_user_id:u1"]}} + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + return 1 + + results = await asyncio.gather(one_request("call-a"), one_request("call-b")) + await asyncio.sleep(0) + assert results == [1, 1] + + +# --------------------------------------------------------------------------- +# Accounting: tokens/dollars via async_log_success_event +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dollar_limit_accounts_usage_and_rejects_once_over(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [{"name": "daily_spend", "tag_id": "end_user_id", "limit": 10.0, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + 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"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key( + time_controller, monkeypatch +): + """ + kwargs at async_log_success_event time is Logging.model_call_details, not + the flat dict admission sees -- for a plain (non LITELLM_METADATA_ROUTES) + chat completion, kwargs["litellm_params"] carries a "litellm_metadata" key + that is always present but set to None, alongside the real, populated + "metadata" dict. get_metadata_variable_name_from_kwargs only checks key + presence, so it always resolved to "litellm_metadata" here and read no + tags/identity at all, silently dropping every token/dollar/key-hash/alias + accounting for this route shape. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [{"name": "daily_spend", "tag_id": "end_user_id", "limit": 10.0, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = { + "litellm_call_id": "call-1", + "litellm_params": { + "litellm_metadata": None, + "metadata": {"tags": ["end_user_id:u1"], "user_api_key": "hash"}, + }, + "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"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_dollar_limit_respects_apply_to_key_alias_at_accounting_time(time_controller, monkeypatch): + """The entry only applies to `premium-key`; a non-listed key's spend must + not be charged against this bucket at all.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "daily_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="other-key"), cache=DualCache(), data=data, call_type="completion" + ) + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"], "user_api_key_alias": "other-key"}, + "standard_logging_object": {"total_tokens": 0, "response_cost": 999.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # premium-key was never charged -- still fully under its own limit. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +# --------------------------------------------------------------------------- +# Config hot-reload: identity-based re-validation, no restart needed +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_config_reload_takes_effect_on_next_request(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 100, "period_seconds": 86400}] + } + }, + ) + 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_type="completion" + ) + + # Reload to a stricter config -- a fresh dict object, matching how a + # proxy config reload replaces litellm_settings.global_tag_rate_limits + # wholesale via setattr(litellm, key, value). A changed `limit` folds + # into the bucket's own policy fingerprint, so this is a fresh counter; + # the new, stricter limit=1 is still reachable in exactly one more call. + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + 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-3"), + 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 883117c223b..f6f0909e1a8 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5806,6 +5806,65 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == primary_model + @pytest.mark.asyncio + async def test_cross_model_scoped_rejection_is_not_retried_via_fallback(self): + """ + veria-ai finding on PR #36541: an entry using ``apply_to_models`` to cap + an entire fallback chain as one unit is defeated by this exact mechanism + if a fallback model isn't also listed in ``apply_to_models`` -- the + rejection here is a deliberate "this whole chain is capped" decision, + not a "this one model is unhealthy" signal, so retrying against an + unlisted fallback silently serves a request the operator's policy meant + to block. ``detail["cross_model_scope"]`` is the marker + global_tag_rate_limits_hook sets for exactly this case; the fallback + handler must re-raise immediately instead of trying any fallback model. + """ + 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}) + + call_count = 0 + + async def mock_pre_call_logic(**kwargs): + nonlocal call_count + call_count += 1 + raise ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "cross_model_scope": True}, + headers={"retry-after": "30"}, + ) + + mock_router = MagicMock() + mock_router.fallbacks = [{"opus-chain": ["sonnet-chain"]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + with pytest.raises(ProxyRateLimitError): + 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 call_count == 1 + assert processor.data["model"] == primary_model + @pytest.mark.asyncio async def test_real_parallel_request_limiter_model_tpm_limit_triggers_fallback(self): """