diff --git a/litellm/__init__.py b/litellm/__init__.py index ec2960c196e..b82c918230d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -119,6 +119,8 @@ _custom_logger_compatible_callbacks_literal = Literal[ "litellm_agent", "dynamic_rate_limiter", "dynamic_rate_limiter_v3", + "model_based_tag_rate_limits_hook", + "global_tag_rate_limits_hook", "langsmith", "prometheus", "otel", @@ -391,6 +393,9 @@ cache: Optional["Cache"] = None # cache object <- use this - https://docs.litel 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/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 38a9966f9f9..284e780cc6f 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -163,7 +163,12 @@ class InMemoryCache(BaseCache): return self.cache_dict[key] = value - if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl + # refresh_ttl bypasses allow_ttl_override's "leave a still-live ttl + # alone" guard -- a caller only sets it for a counter whose ttl must + # keep extending on every write (e.g. a concurrency reservation's + # crash-safety-net ttl), never for one that must stay fixed to its + # original epoch window (e.g. a fixed-period rate-limit bucket). + if kwargs.get("refresh_ttl") or self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) @@ -248,7 +253,15 @@ class InMemoryCache(BaseCache): ) -> list[float] | None: results: Final = [] for increment in increment_list: - result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) + # Each operation's own ttl must reach set_cache, or a key with no + # live ttl yet falls through to the cache's short default_ttl + # instead of the caller's real (often much longer) window -- + # allow_ttl_override already leaves an existing, still-live ttl + # untouched on a later increment, so passing this through on + # every call is safe: it only ever takes effect the first time. + result = await self.async_increment( + increment["key"], increment["increment_value"], ttl=increment.get("ttl"), **kwargs + ) results.append(result) return results diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 41caf732db0..139fc2b2e65 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -187,6 +187,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): pass + async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + """ + Called when a client disconnects mid-request, for a callback that reserved + per-request state outside async_log_success_event/async_log_failure_event + (e.g. a concurrency slot admitted before the first response chunk) -- those + two callbacks never run for a client disconnect, so a callback relying on + them alone to release such state would otherwise leak it until its own + safety-net TTL. + """ + async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"): """Called when an audit log is created. Override in subclasses to handle.""" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3018f0c4d24..4003b5f4d9a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4476,6 +4476,43 @@ def _init_custom_logger_compatible_class( dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router) _in_memory_loggers.append(dynamic_rate_limiter_obj_v3) return dynamic_rate_limiter_obj_v3 + elif logging_integration == "model_based_tag_rate_limits_hook": + from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( + _PROXY_ModelBasedTagRateLimitsHook, # pyright: ignore[reportPrivateUsage] # resolved by name like every other opt-in callback here + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_ModelBasedTagRateLimitsHook): + return callback + + if internal_usage_cache is None: + raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}") + + model_based_tag_rate_limits_hook_obj: Final = _PROXY_ModelBasedTagRateLimitsHook( + internal_usage_cache=internal_usage_cache + ) + + if llm_router is not None and isinstance(llm_router, litellm.Router): + 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") @@ -4916,6 +4953,24 @@ def get_custom_logger_compatible_class( if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): return callback + elif logging_integration == "model_based_tag_rate_limits_hook": + from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( + _PROXY_ModelBasedTagRateLimitsHook, # pyright: ignore[reportPrivateUsage] # resolved by name like every other opt-in callback here + ) + + for callback in _in_memory_loggers: + 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 315fbcba310..f98ff8e8e62 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -33,6 +33,7 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -397,6 +398,32 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return True +async def _release_disconnect_state_on_all_callbacks(request_data: Mapping[str, object]) -> None: + """ + A client disconnect throws GeneratorExit/CancelledError into the streaming + generator, so neither the success nor failure logging callback runs for it + (see the callers of this function). A callback that reserves per-request + state outside of those two callbacks (e.g. a concurrency slot admitted + before the first chunk) would otherwise leak that state until its own + safety-net TTL. Give every registered callback a chance to release such + state via the optional, default-no-op ``async_release_disconnect_state_hook``. + + Only ``CustomLogger`` instances are considered, never raw string entries: + by the time a request can reach this proxy-only cleanup path, startup's + ``ProxyLogging._init_litellm_callbacks`` has already replaced every string + entry in ``litellm.callbacks`` with its initialized instance in place. + """ + for callback in litellm.callbacks: + if not isinstance(callback, CustomLogger): + continue + try: + await callback.async_release_disconnect_state_hook(request_data) + except Exception as e: # noqa: BLE001 # one callback's cleanup must never block another's or the response teardown + verbose_proxy_logger.debug( + "Failed to run async_release_disconnect_state_hook for %s: %s", type(callback).__name__, e + ) + + async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None: pending_tasks: Final = [task for task in tasks if not task.done()] for task in pending_tasks: @@ -1454,6 +1481,7 @@ async def _cancel_llm_call_on_client_disconnect( async def _await_llm_call_cancelling_on_disconnect( request: Request, llm_api_call: "asyncio.Future[_LlmCallT]", + request_data: Mapping[str, object], ) -> _LlmCallT: disconnect_event: Final = asyncio.Event() monitor: Final = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event)) @@ -1461,6 +1489,14 @@ async def _await_llm_call_cancelling_on_disconnect( return await llm_api_call except asyncio.CancelledError: if disconnect_event.is_set(): + # This cancellation never reaches litellm.utils.wrapper_async's own + # except block (asyncio.CancelledError is a BaseException, not an + # Exception, since Python 3.8), so async_log_failure_event never + # fires for it -- the same gap async_release_disconnect_state_hook + # was added for on the streaming path (see + # _finalize_streaming_generator_cleanup), just reached here via a + # cancelled non-streaming call instead of a mid-stream disconnect. + await _release_disconnect_state_on_all_callbacks(request_data) raise HTTPException( status_code=499, detail=_CLIENT_DISCONNECT_DETAIL, @@ -1967,7 +2003,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( @@ -2006,7 +2045,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 @@ -2285,7 +2336,9 @@ class ProxyBaseLLMRequestProcessing: try: if general_settings.get("cancel_on_disconnect", False): - responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) + responses = await _await_llm_call_cancelling_on_disconnect( # rebind-ok: assigned in exactly one of these two mutually exclusive branches + request, llm_responses, self.data + ) else: responses = await llm_responses finally: @@ -3364,6 +3417,8 @@ class ProxyBaseLLMRequestProcessing: and user_api_key_dict is not None ): await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + if not success_event_owns_slot_release: + await _release_disconnect_state_on_all_callbacks(request_data) if hasattr(response, "aclose"): try: 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..43d2a1462f2 --- /dev/null +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -0,0 +1,791 @@ +""" +Tag-scoped token, request, dollar, and concurrency rate limits declared once, +globally, in `litellm_settings.global_tag_rate_limits` and enforced 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. + +Model-independent sibling of `model_based_tag_rate_limits_hook`, which +enforces the same `TagRateLimitEntry` shape per-deployment instead. Reuses +that sibling's shared helpers (`entry_applies`, the atomic Lua scripts, cache +partitioning, bucket-key hashing) from `tag_rate_limits_shared.py`, but has +its own smaller admission/accounting engine with no routing-group or +per-deployment concerns. + +Three entry-level knobs: `apply_to_key_alias` scopes an entry to specific +virtual-key aliases; `apply_to_models` scopes it to specific caller-facing +model names, letting one entry cap a whole fallback chain as a unit (a +rejection then carries `detail["cross_model_scope"] = True` so +`_pre_call_with_fallbacks` re-raises instead of silently admitting the +request through an unlisted fallback model); `scope_by_key_hash` controls +whether matching keys share one bucket or each gets its own. + +`data["litellm_logging_obj"]` does exist by the time `async_pre_call_hook` +runs, but `_pre_call_with_fallbacks` re-runs the whole pre-call pipeline +(building a fresh `Logging` object each time) once per fallback model on the +same `litellm_call_id`, so unlike `model_based_tag_rate_limits_hook`'s +per-Router-hop reservations, a stash keyed to one attempt's own +`model_call_details` wouldn't survive to a later attempt. Per-request state +instead lives on a `ContextVar`-based stash (the same pattern +`parallel_request_limiter_v3.py` uses for the identical admission-to-release +problem), keyed by `litellm_call_id` so a nested LiteLLM call sharing the +same inherited context (an LLM-judge guardrail, for example) gets its own +isolated entry instead of releasing the outer call's still-pending +reservation early. Confirmed live: a real streaming client disconnect +correctly releases its concurrency reservation through this mechanism. +""" + +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.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] # shared private helper, reused by model_based_tag_rate_limits_hook too +) +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 ( + 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, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + policy_fingerprint as _policy_fingerprint, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + resolve_authoritative_metadata_variable_name as _resolve_authoritative_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] # shared private helper, reused by model_based_tag_rate_limits_hook too +) +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import TagRateLimitEntry, TagRateLimits + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span: TypeAlias = _Span +else: + Span: TypeAlias = object + + +def _entry_applies_any_admitted_model( + entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, admitted_models: frozenset[str] +) -> bool: + """Same as `_entry_applies`, except an `apply_to_models`-scoped entry + counts as applying if ANY model an admission attempt for this call_id + saw was in scope -- not just whichever model the call ultimately served. + A `_pre_call_with_fallbacks` retry re-admits with a different model for + the same call_id, and an entry that matched an earlier attempt must + still get its success-time accounting.""" + if not admitted_models: + return _entry_applies(entry, tags, key_alias, None) + return any(_entry_applies(entry, tags, key_alias, model) for model in admitted_models) + + +def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: + """Namespaced under `tag_rl:global:` so it never collides with + `model_based_tag_rate_limits_hook`'s own `tag_rl:{model_group}:...` keys.""" + 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` rather than one shared mutable instance so a + nested LiteLLM call (an LLM-judge guardrail) that mints its own call id + but inherits the same context doesn't release the outer call's + still-pending reservation early. + """ + + admission_time: float | None = None + # Every model any admission attempt for this call_id has classified + # entries against, accumulated rather than overwritten: a + # _pre_call_with_fallbacks retry re-runs admission with a *different* + # model for the same call_id, and an apply_to_models entry that matched + # an earlier attempt must still get its accounting at success time even + # though the request ultimately serves from a later attempt's model. + admitted_models: frozenset[str] = field(default_factory=frozenset) + pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue + # Keys already charged for this call_id, so a fallback retry (same + # litellm_call_id, different model) renews instead of double-charging. + charged_request_keys: list[str] = field(default_factory=list) # mutable-ok: see comment above + # key_hash of whoever first claimed this stash. litellm_call_id is + # caller-controlled (x-litellm-call-id), so only a later admission with + # the same authenticated key_hash may renew this stash's charges. + 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, refresh_ttl: bool + ) -> 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, 1 if refresh_ttl else 0) + ) + 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, refresh_ttl=refresh_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, bool]], + ) -> tuple[int | None, tuple[float, ...]]: + """All-or-nothing atomic admission across `checks`: on a rejection, + refunds every check admitted earlier in this batch.""" + if not checks: + return None, () + admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection + for index, (cache, key, limit, increment, ttl, refresh_ttl) in enumerate(checks): + admitted = False + try: + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl, refresh_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, bool]], up_to_index: int + ) -> None: + for refund_index in range(up_to_index): + refund_cache, refund_key, _limit, refund_increment, _ttl, _refresh_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 _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": + 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 deliberately excluded: it can resolve from + # inherited_tags (server-assigned key/team/project metadata), + # and echoing it back would disclose that identity to the + # caller. verbose_proxy_logger.debug above still logs it + # server-side for observability. + "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 + + # _pre_call_with_fallbacks can re-run this pipeline once per fallback + # model on any ProxyRateLimitError, reusing the same litellm_call_id -- + # see charged_request_keys for how a repeat run renews instead of + # re-charging. + stash: Final = _claim_stash_for_data(data) + + # Not get_metadata_variable_name_from_kwargs (naive key-presence + # check): a caller can forge an empty (or None) litellm_metadata on + # an ordinary request to make that check pick it over the real, + # populated metadata the proxy wrote identity/tags into, seeing no + # tags at all and admitting past every configured limit. + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(data) + 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 + + # First admission for this stash claims ownership; only a later one + # with the same key_hash may renew its charges (see owner_key_hash). + 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 + 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) + 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 fallback attempt for the same request) renews + # at zero net cost instead of charging a second unit. + 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), + check.unit == "concurrency", + ) + 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] + ) + + # Exclude already_reserved_concurrency_keys: that key renewed at + # zero cost above, so re-adding it would make release decrement + # twice for a counter 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, so a call_id collision from + # a different key_hash can't contaminate the rightful owner's + # renewal tracking. + 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 + + self._record_admitted_model(stash, model, renewal_allowed) + return data + + async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None: + stash: Final = _stash_for_call(_call_id_from_kwargs(request_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_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + await self._release_pending_for_call_id(request_data) + + async def async_post_call_failure_hook( + self, + request_data: dict, # mutable-ok: must match CustomLogger.async_post_call_failure_hook's own base signature exactly + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + """ + A request that never reaches Router (every fallback model also + rejected, or none configured) never runs the actual LLM call, so + neither async_log_success_event nor async_log_failure_event -- both + tied to that call's own wrapper -- ever fires for it. This is the + only remaining release path for a reservation from an earlier, + successful admission attempt in the same _pre_call_with_fallbacks + chain. litellm_call_id survives proxy/utils.py's own stripping here + (only litellm_logging_obj is popped), so the same ContextVar-based + stash lookup as the other release hooks still works. + """ + await self._release_pending_for_call_id(request_data) + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime | None, + end_time: datetime | None, + ) -> None: + # Always release regardless of which hook raised: this hook's own + # rejection never reserves a slot, so pending_concurrency_keys is + # already empty in that case and the check below no-ops; a rejection + # from model_based_tag_rate_limits_hook (same error marker) can still + # land after this hook already reserved its own slot. + await self._release_pending_for_call_id(kwargs) + + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime | None, + end_time: datetime | None, + ) -> 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 = kwargs.get("standard_logging_object") + if not isinstance(standard_logging_object, dict): + 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_raw: Final = kwargs.get("litellm_params") + litellm_params_for_metadata: Final[Mapping[str, object]] = ( + litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs + ) + metadata_variable_name: Final = _resolve_authoritative_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 = _order_tags_for_identity_resolution( + _get_tags_from_request_kwargs(litellm_params_for_metadata, metadata_variable_name=metadata_variable_name), + litellm_params_for_metadata, + 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() + ) + admitted_models: Final = stash.admitted_models if stash is not None else frozenset() + 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_any_admitted_model(entry, tags, key_alias, admitted_models): + 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=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/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py new file mode 100644 index 00000000000..b425571d0c2 --- /dev/null +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -0,0 +1,1916 @@ +""" +Tag-scoped token, request, dollar, and concurrency rate limits, admitted +once per routing hop via `async_filter_deployments`, for limits declared +per-deployment under `model_info.tag_rate_limits`. + +Shares its identity/scope-extraction, policy-fingerprinting, bucket-key +hashing, and cache-partitioning primitives with the model-independent +`global_tag_rate_limits_hook.py` via `tag_rate_limits_shared.py`; this +module owns everything specific to per-deployment admission instead: +routing-group resolution, the (team-alias-aware) limits index, and +per-deployment dedup. +""" + +import asyncio +import contextvars +import json +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime +from itertools import groupby +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, NamedTuple, 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_parent_otel_span_from_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching dynamic_rate_limiter_v3's identical import +) +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] # this hook explicitly reuses its Redis/TTL-preserving increment machinery, see module docstring +) +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 ( + EMPTY_MAPPING as _EMPTY_MAPPING, +) +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 ( + ScopeSignature as _ScopeSignature, +) +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 ( + 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, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + policy_fingerprint as _policy_fingerprint, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + resolve_authoritative_metadata_variable_name as _resolve_authoritative_metadata_variable_name, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + scope_signature as _scope_signature, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.router import Router +from litellm.router_strategy.tag_based_routing import ( + _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching dynamic_rate_limiter_v3's identical import +) +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.llms.openai import AllMessageValues +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 + +# (tag_id, name, limit, period_seconds, scope_by_key_hash, enabled_for, +# disabled_for, apply_to_key_alias, apply_to_models) -- the fields that +# decide whether two deployments' entries are the same rate limit for dedup +# purposes; see _build_group_limits. Two deployments that agree on the first +# five but disagree on any scoping field are declaring genuinely different +# policies (e.g. one excludes a user the other doesn't) and must not be +# merged into one shared bucket -- the same class of bug this signature +# already guards against for a plain divergent `limit`. +_DedupSignature: TypeAlias = tuple[ + str, + str, + float, + int, + bool, + _ScopeSignature, + _ScopeSignature, + tuple[str, ...] | None, + tuple[str, ...] | None, +] + + +@dataclass(frozen=True) +class _ConfiguredLimit: + unit: _LimitUnit + entry: TagRateLimitEntry + # None => chain-wide (every deployment in the model_group shares one + # bucket). Otherwise the sorted deployment ids that declared this exact + # value -- the bucket is shared among only those deployments. + deployment_scope: tuple[str, ...] | None + # The owning team_id, set whenever this limit belongs to a team-owned + # deployment -- via `by_team_alias`, or via `by_model_name` for a group + # any of whose deployments also declare a `team_public_model_name` (see + # `_build_limits_index`). team_public_model_name is only unique per team, + # so two teams can publish the identical alias string; without the + # team_id folded into the bucket key too, both teams' identically-named, + # identically-configured limits would collide on the same Redis counter + # despite the index itself correctly scoping the `by_team_alias` lookup + # by (team_id, alias). + team_scope: str | None = None + # Overrides the caller-visible name in the bucket key whenever that name + # doesn't uniquely identify the bucket. Two independent cases set this: + # (1) `resolve()`'s direct lookup missed and `resolve_any()` fell back to + # resolving via a candidate deployment's own model_name instead (routing + # groups, and any other indirection where Router deliberately keeps the + # caller-visible name distinct from every deployment's own model_name) -- + # here it's stamped with that deployment's own model_name, so limits from + # two different underlying model_names sharing one routing group never + # collide on one counter; (2) the limit belongs to a team-owned + # deployment (`team_scope` is set) -- here it's stamped with the team's + # own `team_public_model_name`, the same value regardless of whether this + # entry was found via `by_team_alias` (caller used the public alias) or + # via `by_model_name` (caller used the deployment's own internal, + # auto-generated model_name), so both call shapes land on one shared + # bucket instead of splitting a team's usage across two counters. None + # when resolved directly and not team-owned, in which case the + # caller-visible name is already unambiguous and safe to hash by. + resolved_group: str | None = None + + +def _deployment_id(deployment: Mapping[str, object]) -> str | None: + return (deployment.get("model_info") or _EMPTY_MAPPING).get("id") + + +def _model_name_of(deployment: Mapping[str, object]) -> str: + """`model_name` is a required field on every deployment dict; `str(...)` + (rather than a bare index) gives `sorted`/`groupby`'s key functions a + provably orderable, hashable return type without an unsafe cast -- every + real deployment's value here is already a string, so this is a no-op + coercion in practice.""" + return str(deployment["model_name"]) + + +def _extract_team_id(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Reads `user_api_key_team_id` from only the one field + `_resolve_authoritative_metadata_variable_name` names as authoritative for + this request -- never falling back to the other field, since + `litellm_pre_call_utils.py` writes the real, server-authenticated value + into that one field alone and leaves the other exactly as the caller + sent it. An OR-fallback across both would let a caller's own + `metadata.user_api_key_team_id` (still present, unvalidated, on a route + where `litellm_metadata` is the authoritative field) win over the real + value.""" + active: Final = request_kwargs.get(metadata_variable_name) or _EMPTY_MAPPING + team_id: Final = active.get("user_api_key_team_id") + return team_id if isinstance(team_id, str) else None + + +def _entries_for_unit(deployment: Mapping[str, object], unit: _LimitUnit) -> tuple[TagRateLimitEntry, ...]: + raw_tag_rate_limits: Final = (deployment.get("model_info") or _EMPTY_MAPPING).get("tag_rate_limits") + if not raw_tag_rate_limits: + return () + tag_rate_limits: Final = TagRateLimits.model_validate(raw_tag_rate_limits) + group: Final = getattr(tag_rate_limits, _UNIT_TO_GROUP_FIELD[unit]) + return tuple(group.limits) if group is not None else () + + +def _configured_limit_for_signature( + unit: _LimitUnit, + entry: TagRateLimitEntry, + declaring_ids: Sequence[str], + is_chain_wide: bool, +) -> _ConfiguredLimit | None: + if unit == "concurrency" and not is_chain_wide: + verbose_proxy_logger.warning( + "model_based_tag_rate_limits_hook: concurrency_limits entry %r (tag_id=%s) is not declared identically by every " + "deployment sharing this model_name; per-deployment-scoped concurrency limits are not supported " + "and this entry is being skipped entirely.", + entry.name, + entry.tag_id, + ) + return None + return _ConfiguredLimit( + unit=unit, + entry=entry, + deployment_scope=None if is_chain_wide else tuple(sorted(declaring_ids)), + ) + + +def _build_group_limits(deployments: Sequence[Mapping[str, object]], unit: _LimitUnit) -> tuple[_ConfiguredLimit, ...]: + """ + One `_ConfiguredLimit` per distinct (tag_id, name, limit, period_seconds) + declared for `unit` across `deployments` (all sharing one `model_name`). + + A signature declared identically by every deployment in the group is + chain-wide (one shared bucket, regardless of which deployment serves). + A signature declared by only some deployments, or where deployments + genuinely disagree on the value for the same (tag_id, name), becomes a + per-deployment-scoped bucket shared by exactly the deployments that + declared that value -- silently dropping a divergent deployment's config + (as a naive dedupe-by-name index would) is the exact bug this guards + against. + + `concurrency` is the one exception: a per-deployment-scoped reservation + is never created for it. Admission for a hop reserves every scope whose + deployments overlap `healthy_deployments`, but only one deployment ends + up actually serving -- releasing the exact reservation(s) that were never + used, without a per-request slot identity to track which reservation + belongs to which hop, isn't solved correctly by this design (a + since-fixed live bug: an admitted-then-failed call's per-deployment + reservation was never released; a caller could also strand a sibling + deployment's reservation just by never being routed to it). A divergent + concurrency signature is dropped with a warning instead of silently + creating a bucket that can leak; only chain-wide concurrency entries + (identical across every deployment in the group) are supported. + """ + # Insertion order here is load-bearing: it decides which limit's + # ProxyRateLimitError surfaces first when several are breached by the + # same hop (see async_filter_deployments). A presort-based + # itertools.groupby would need to sort by signature to group it, which + # would scramble that first-seen order, so this stays a plain + # accumulator instead. + declaring_ids_by_signature: Final = {} # mutable-ok: first-seen order here decides which limit's error raises first (see comment above); sorting to use groupby would scramble it + # The dedup signature is deliberately narrower than the full entry: two + # deployments agreeing on (tag_id, name, limit, period_seconds, + # scope_by_key_hash) share one bucket even if they set key_ttl_seconds or + # max_in_memory_cache_size differently. Whichever deployment's entry is + # seen first for a given signature supplies those fields for the whole + # group -- an arbitrary but deterministic tie-break, consistent with the + # first-seen-order precedent already established above. + representative_entry_by_signature: Final[dict[_DedupSignature, TagRateLimitEntry]] = {} # mutable-ok: see above + for deployment in deployments: + dep_id = _deployment_id(deployment) + if dep_id is None: + continue + for entry in _entries_for_unit(deployment, unit): + signature = ( + entry.tag_id, + entry.name, + entry.limit, + entry.period_seconds, + entry.scope_by_key_hash, + _scope_signature(entry.enabled_for), + _scope_signature(entry.disabled_for), + entry.apply_to_key_alias, + entry.apply_to_models, + ) + ids_for_signature = declaring_ids_by_signature.setdefault(signature, []) # mutable-ok: see comment above + # One deployment declaring the identical entry twice (a config + # duplicate) must count once, or len(declaring_ids) inflates past + # total_deployments below, making is_chain_wide false for an + # entry every deployment actually agrees on -- for concurrency + # that silently drops the entry entirely (see the docstring + # above), disabling enforcement rather than degrading it. + if dep_id not in ids_for_signature: + ids_for_signature.append(dep_id) # mutable-ok: see comment above + representative_entry_by_signature.setdefault(signature, entry) # mutable-ok: see comment above + + distinct_signature_count_by_name: Final[Mapping[tuple[str, str], int]] = MappingProxyType( + { + (tag_id, name): sum( + 1 + for other_tag_id, other_name, *_rest in declaring_ids_by_signature + if (other_tag_id, other_name) == (tag_id, name) + ) + for tag_id, name, *_rest in declaring_ids_by_signature + } + ) + + total_deployments: Final = len(deployments) + configured: Final = tuple( + configured_limit + for signature, declaring_ids in declaring_ids_by_signature.items() + if ( + configured_limit := _configured_limit_for_signature( + unit, + representative_entry_by_signature[signature], + declaring_ids, + is_chain_wide=( + distinct_signature_count_by_name[(signature[0], signature[1])] == 1 + and len(declaring_ids) == total_deployments + ), + ) + ) + is not None + ) + return configured + + +@dataclass(frozen=True, slots=True) +class _LimitsIndex: + """ + Two lookup tables because a `team_public_model_name` alias is only + unique per team, not globally: Router itself lets different teams + publish the identical alias string for different deployments, and + resolves each caller's own team's deployment by `(team_id, name)`, not + by `name` alone (see `Router._update_team_model_index`). Keying alias + limits by name alone here would let one team's config silently + overwrite another's. + """ + + by_model_name: Mapping[str, tuple[_ConfiguredLimit, ...]] + by_team_alias: Mapping[tuple[str, str], tuple[_ConfiguredLimit, ...]] + + def resolve(self, model: str, team_id: str | None) -> tuple[_ConfiguredLimit, ...]: + if team_id is not None: + scoped: Final = self.by_team_alias.get((team_id, model)) + if scoped is not None: + return scoped + return self.by_model_name.get(model, ()) + + def resolve_any( + self, model: str, team_id: str | None, candidate_model_names: Iterable[str] + ) -> tuple[_ConfiguredLimit, ...]: + """ + Like `resolve`, but falls back to each candidate deployment's own + `model_name` when `model` itself matches neither table -- Router + deliberately keeps `model` as a callable routing-group name distinct + from every member deployment's own `model_name` (see + `Router._get_routing_group_deployments`), so a group-addressed call + would otherwise never match this index at all despite its member + deployments carrying real `tag_rate_limits`. Each fallback result is + stamped with the `model_name` it actually came from (`resolved_group`) + so hashing stays namespaced per underlying group even when the + candidates span more than one `model_name`. + + Members declaring the identical signature and scope are deduped to + one shared entry: only one deployment in the group ends up actually + serving a given hop, but every member's own `model_name` is resolved + independently above, so an undeduped union would check and charge + every member's bucket for that one hop -- request/concurrency + capacity a caller never actually used, and a false 429 for a sibling + member that was never over its own limit. Divergent configs across + model_names (different limit/period/scope for the same tag_id+name) + are left as separate entries, same as before this dedup: resolving + that ambiguity needs knowing which deployment will be picked, which + isn't known yet at this admission-time hook. + + Candidates are deduped in sorted order, not raw `frozenset` iteration + order: `frozenset` order depends on the process's hash seed, so two + workers resolving the identical candidate set could otherwise pick + different members as `resolved_group` and end up checking/accounting + against different Redis keys for what's meant to be one shared bucket. + + A candidate's own entry can already carry `resolved_group` set to its + team's `team_public_model_name` (see `_build_limits_index`) -- that + stamp is left untouched rather than overwritten with `name` here, or + a team-owned deployment reachable through *both* its team alias and a + routing group would hash to two different buckets depending on which + path a caller happened to take. + """ + direct: Final = self.resolve(model, team_id) + if direct: + return direct + deduped: Final[dict[tuple[object, ...], _ConfiguredLimit]] = {} # mutable-ok: see docstring above + for name in sorted(frozenset(candidate_model_names)): + for limit in self.by_model_name.get(name, ()): + key = ( + limit.unit, + limit.entry.tag_id, + limit.entry.name, + limit.entry.limit, + limit.entry.period_seconds, + limit.entry.scope_by_key_hash, + _scope_signature(limit.entry.enabled_for), + _scope_signature(limit.entry.disabled_for), + limit.entry.apply_to_key_alias, + limit.entry.apply_to_models, + limit.deployment_scope, + limit.team_scope, + ) + deduped.setdefault( # mutable-ok: see docstring above + key, limit if limit.resolved_group is not None else replace(limit, resolved_group=name) + ) + return tuple(deduped.values()) + + +def _team_alias_key(deployment: Mapping[str, object]) -> tuple[str, str] | None: + model_info: Final = deployment.get("model_info") or _EMPTY_MAPPING + team_id: Final = model_info.get("team_id") + team_public_model_name: Final = model_info.get("team_public_model_name") + if team_id and team_public_model_name: + return (team_id, team_public_model_name) + return None + + +def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIndex: + """ + `by_model_name` is keyed by every deployment's own `model_name`, grouping + deployments that share one. + + `by_team_alias` additionally covers `team_public_model_name`: a team + calling through its own public alias reaches `async_filter_deployments` + with that alias as `model`, while the deployment dicts in + `healthy_deployments` still carry their own real `model_name` -- + `Router` never rewrites it for this path (unlike `model_group_alias`, + which is resolved to the real model_name before routing even starts). + Without this, tag limits configured on a team-aliased chain would never + be looked up at all. + + This is a genuinely separate grouping from `by_model_name`, not a lookup + into it: litellm auto-generates each team-added deployment's own + `model_name` as `model_name_{team_id}_{uuid}` (see + `model_listing_utils.py`), so multiple deployments sharing one + `team_public_model_name` alias routinely have different, unique + `model_name` values -- Router's own `team_model_to_deployment_indices` + aggregates them by `(team_id, team_public_model_name)` regardless. + Computing alias limits once per `model_name` group and keying the alias + to whichever group happened to declare it would drop every other + same-alias group's limits whenever more than one model_name shares an + alias, since the last one processed would silently overwrite the rest. + `_build_group_limits` has no `model_name`-specific logic (it only reads + each deployment's own id and its own entries), so it's safe to reuse + unchanged for a deployment set spanning multiple `model_name` values. + + Deployments are grouped via a stable sort + itertools.groupby rather than + a setdefault-in-a-loop accumulator: `sorted` is stable, so deployments + sharing a key keep the exact same relative order `_build_group_limits` + would have seen them in without the sort, which is what keeps this safe + (that relative order decides first-seen signature order downstream). + """ + sorted_by_model_name: Final = sorted(model_list, key=_model_name_of) + by_model_name: Final[Mapping[str, tuple[_ConfiguredLimit, ...]]] = MappingProxyType( + { + model_name: ( + # `Router.should_include_deployment` lets a same-team caller + # reach a team-owned deployment by its own internal + # model_name, not only its team_public_model_name alias + # (litellm auto-generates a name unique per (team_id, uuid), + # so every deployment in this group shares one team_id when + # any does). Stamping team_scope alone is not enough to unify + # this with the alias entry below: `_hash_tag` still hashes + # the caller-visible name by default, and that name differs + # between the two paths (the internal model_name here vs. the + # public alias below). Also stamping `resolved_group` with the + # team's own alias forces both paths to hash under the + # identical name, so a caller can't split its usage across two + # independent counters just by alternating which name it calls. + tuple(replace(limit, team_scope=team_key[0], resolved_group=team_key[1]) for limit in configured) + if (team_key := next((key for dep in group if (key := _team_alias_key(dep))), None)) is not None + else configured + ) + for model_name, deployment_group in groupby(sorted_by_model_name, key=_model_name_of) + for group in (tuple(deployment_group),) + if (configured := tuple(limit for unit in _LIMIT_UNITS for limit in _build_group_limits(group, unit))) + } + ) + + aliased: Final = tuple( + (alias_key, deployment) for deployment in model_list if (alias_key := _team_alias_key(deployment)) is not None + ) + sorted_by_alias: Final = sorted(aliased, key=lambda pair: pair[0]) + by_team_alias: Final[Mapping[tuple[str, str], tuple[_ConfiguredLimit, ...]]] = MappingProxyType( + { + alias_key: alias_configured + for alias_key, alias_group in groupby(sorted_by_alias, key=lambda pair: pair[0]) + for aliased_group in (tuple(dep for _key, dep in alias_group),) + if ( + alias_configured := tuple( + # resolved_group is already the caller-visible name on + # this path (the caller reached this group by dialing the + # alias directly), but stamping it explicitly keeps both + # index branches symmetric and independent of whatever + # value the caller happens to pass as `model_group` into + # `_hash_tag`. + replace(limit, team_scope=alias_key[0], resolved_group=alias_key[1]) + for unit in _LIMIT_UNITS + for limit in _build_group_limits(aliased_group, unit) + ) + ) + } + ) + + return _LimitsIndex(by_model_name=by_model_name, by_team_alias=by_team_alias) + + +# Upper bound on how stale the limits index may be after a length-preserving +# deployment update (e.g. editing an existing deployment's tag_rate_limits +# in place via the admin API, which never changes len(model_list)). Router +# exposes no generic "config changed" version counter to key off instead, so +# this bounds staleness by simply re-checking periodically. +_INDEX_TTL_SECONDS: Final = 5.0 + + +# Concurrency reservation keys accumulated for the current logical request, +# not yet released, paired with the cache-size override (from +# TagRateLimitEntry.max_in_memory_cache_size) each reservation was +# incremented under: releasing a reservation must decrement the exact same +# cache partition it was incremented on, or the release silently no-ops on +# the wrong (default) partition and the reservation leaks forever. +# +# Stashed directly on `Logging.model_call_details` under this field, not a +# `contextvars.ContextVar`: the real proxy request pipeline forks the +# streaming response through several distinct asyncio Tasks (the disconnect +# race in `create_response`, the streaming generator's own task, ...), and a +# ContextVar only propagates forward into tasks forked *after* a value was +# `.set()` -- a task that isn't a descendant of admission's task never sees +# it, so release silently finds nothing and every reservation leaks until +# `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, disconnect or not (confirmed live: +# even a fully-completed, non-disconnected streaming request never released +# its slot). `model_call_details` is a single dict, explicitly passed by +# object reference through both admission's `request_kwargs` (as +# `request_kwargs["litellm_logging_obj"].model_call_details`) and release's +# `kwargs` (`async_log_success_event`/`async_log_failure_event`'s `kwargs` +# argument *is* `model_call_details` -- see their own callers), so it +# survives task boundaries by construction, not by ambient context. +# +# Deliberately not keyed by `litellm_call_id` instead: that field is +# caller-controlled via the `x-litellm-call-id` request header, so two +# unrelated concurrent requests sharing a caller-chosen id would merge their +# reservations under a shared identifier -- letting one request's release +# free a different request's still-live slot. `model_call_details` is a +# plain Python object with no caller-visible identifier, created fresh +# server-side per logical request (and shared across that request's own +# fallback hops, matching the original chain-wide release semantics), so it +# can't be forged or guessed. +# +# "shared across that request's own fallback hops" is not the same as +# "shared across every task that happens to touch model_call_details": +# `Router.abatch_completion`'s comma-separated multi-model dispatch runs +# several branches concurrently, each its own Task, but hands every branch +# the identical `litellm_logging_obj` -- so a still-live sibling branch's +# own reservation can sit in this same list. Each entry also carries an +# admission-scoped token (see `_current_admission_token`), but that token is +# only ever trusted at the *next admission's own* stale-hop cleanup (see +# `_release_stale_hop_reservations`), never at the terminal release hooks +# below: a `ContextVar` has the identical non-descendant-task blind spot +# documented above for the reservation data itself, so a streaming +# response's own success event -- fired from a task the proxy forked +# independently of admission's -- would never see a matching token either, +# leaking every streaming request's reservation instead of releasing it. +# The terminal hooks release unconditionally, accepting the narrower risk +# that a still-live `abatch_completion` sibling's own reservation gets +# freed early, in exchange for actually releasing the far more common +# single-branch (including streaming) case. +_PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys" + +# Identifies which admission call queued a given reservation, scoped by +# asyncio Context rather than by Task identity: `asyncio.create_task` +# snapshots the *current* Context into the new Task, so a task explicitly +# forked from within one hop's own admission (e.g. litellm's own +# `LoggingWorker.enqueue` dispatch, which explicitly propagates the calling +# context) still reads back the same token, while `abatch_completion`'s +# sibling branches -- forked *before* any of them ever called admission -- +# each start from an unset ContextVar and mint their own distinct token on +# first use, never matching each other's. Only safe where the reader is +# guaranteed to run in a descendant of the writer's task -- see +# `_PENDING_CONCURRENCY_KEYS_FIELD`'s own docstring for where that doesn't +# hold. +_ADMISSION_CONTEXT: Final[contextvars.ContextVar[object | None]] = contextvars.ContextVar( + "_model_based_tag_rate_limits_admission_context", default=None +) + + +def _current_admission_token() -> object: + token: Final = _ADMISSION_CONTEXT.get() + if token is not None: + return token + fresh_token: Final = object() + _ADMISSION_CONTEXT.set(fresh_token) + return fresh_token + + +# Same `model_call_details`-stashing rationale as the field above, for a +# different unit: "requests" is atomic and admitted once per hop (see +# _ATOMIC_UNITS), same as concurrency, but a "requests" limit is meant to cap +# logical client requests, not internal routing attempts -- a chain that +# fails once before succeeding must still consume exactly one unit overall, +# not one per hop. _release_stale_hop_reservations refunds a stale entry +# here the same way it releases a stale concurrency reservation, since its +# own invariant (a queued entry still present when a new hop's admission +# runs can only belong to an earlier hop of this same request that already +# failed) holds identically for either unit. Unlike concurrency, a +# successful (or chain-final-failing) hop's own entry here is deliberately +# never refunded -- exactly one unit must survive per logical request -- so +# async_log_success_event/async_log_failure_event must leave this field +# completely untouched: litellm's has_logged_async_failure dedup lets the +# *first* failing hop's own failure event through (not only a chain's final +# failure), so popping this field there -- even just to discard it -- would +# strand the very entry the *next* hop's admission is relying on being able +# to refund. There is no final-hop/cache-mirror problem to solve for this +# field either: the one hop that never gets superseded is exactly the one +# whose charge should stick, with nothing left to clean up. +_PENDING_REQUEST_INCREMENTS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_request_increments" + +# Mirrors the latest hop's own queued reservation in the same external cache +# the reservations themselves live in, keyed by (litellm_call_id, key_hash), +# for the one release path that cannot reach model_call_details at all: +# proxy/utils.py's post_call_failure_hook deliberately pops litellm_logging_obj +# off request_data before invoking any callback's async_post_call_failure_hook +# ("Remove before callbacks iterate — not serialisable"), so a fallback +# chain's own final, chain-exhausting failure -- which only this hook fires +# for, since litellm's has_logged_async_failure dedup blocks +# async_log_failure_event for every hop after the first -- has no +# model_call_details to pop a reservation from. +# +# Neither a ContextVar nor the flat request_kwargs dict works here (both +# confirmed live, not just reasoned about): a ContextVar's value only +# propagates into descendant tasks, and Router's own per-hop/per-attempt +# execution does not keep the task that later calls post_call_failure_hook +# a descendant of the task that ran the final hop's own admission, so a +# value set there is invisible by the time this fires. request_kwargs is a +# distinct object every hop (confirmed via id()), is a third, unrelated +# object again by the time post_call_failure_hook runs, and mutating it +# directly leaks the mutated key into the actual provider call as an +# `extra_body` param, since litellm forwards unrecognized kwargs verbatim. +# litellm_call_id is the one identifier that is stable across every one of +# those objects, so an external cache keyed by it -- the same Redis/ +# in-memory store the reservations themselves already live in -- is the +# only channel that survives all three failure modes at once. +# +# litellm_call_id alone is not enough to key this cache: it comes from the +# caller-controlled x-litellm-call-id header, so two unrelated requests that +# choose the identical id would overwrite each other's mirror entry, letting +# one caller's terminal failure release a completely different caller's +# still-live reservation. Folding in key_hash -- the calling virtual key's +# hash, resolved server-side (UserAPIKeyAuth.api_key in +# async_post_call_failure_hook, metadata["user_api_key"] everywhere else, +# both authenticated before this hook ever runs) -- confines a collision to +# one caller reusing its own call_id across two of its own concurrent +# requests. This is not confined to "weakens only that caller's own cap": +# for a chain-wide (non scope_by_key_hash) tag, the released reservation is +# on a bucket that tag value's other callers share too, so the forging +# caller's own terminal failure can release a slot on a bucket a different +# caller is also drawing from. Closing this fully needs a per-admission +# identifier that is both server-generated (unlike litellm_call_id) and +# survives proxy/utils.py's post_call_failure_hook stripping +# litellm_logging_obj (unlike everything this comment already ruled out +# above) -- tracked as a known follow-up rather than attempted here. +_PENDING_RESERVATIONS_CACHE_KEY_PREFIX: Final = "model_based_tag_rate_limits:pending_reservations:" + + +def _pending_reservations_cache_key(call_id: str, key_hash: str | None) -> str: + # call_id is caller-controlled (the x-litellm-call-id header) with no + # length bound -- same unbounded-cache-key concern _fixed_length_identity + # documents for tag values, reused here rather than duplicated. + return f"{_PENDING_RESERVATIONS_CACHE_KEY_PREFIX}{_fixed_length_identity(call_id)}:{key_hash or ''}" + + +def _encode_reservations(reservations: Sequence[tuple[str, "_PartitionKey"]]) -> str: + return json.dumps( + tuple( + (key, partition_key if partition_key is None else tuple(partition_key)) + for key, partition_key in reservations + ) + ) + + +def _as_decoded_list(raw: object) -> Sequence[object] | None: + # InMemoryCache.get_cache always attempts json.loads on read regardless + # of what was stored (see its own implementation), so a value written as + # our own already-JSON-encoded string comes back pre-decoded into a list + # when served from the in-memory layer; only a real Redis round trip + # hands back the raw string that still needs decoding here. + if isinstance(raw, list): + return raw + if not isinstance(raw, str): + return None + try: + decoded: Final = json.loads(raw) + except (TypeError, ValueError): + return None + return decoded if isinstance(decoded, list) else None + + +def _decode_reservations(raw: object) -> tuple[tuple[str, "_PartitionKey"], ...]: + decoded: Final = _as_decoded_list(raw) + if decoded is None: + return () + entries: Final = [] # mutable-ok: accumulator over an externally-decoded, untrusted-shape list; immediately frozen below + for item in decoded: + if not (isinstance(item, list) and len(item) == 2 and isinstance(item[0], str)): + continue + partition_key_raw = item[1] + partition_key: _PartitionKey = tuple(partition_key_raw) if isinstance(partition_key_raw, list) else None # pyright: ignore[reportGeneralTypeIssues] # decoded from our own _encode_reservations output; shape validated above + entries.append((item[0], partition_key)) # mutable-ok: see comment above + return tuple(entries) + + +# The admission-time timestamp a hop's token/dollar checks classified their +# bucket against, stashed on the same model_call_details object so success +# accounting recomputes the identical bucket_id (int(now) // period_seconds) +# instead of a fresh one. A completion can take long enough for a fresh +# timestamp at success time to land in the *next* window than the one +# admission actually checked, letting a burst of calls admitted against one +# (still-under-limit) window get charged entirely into the next window's +# fresh, unrelated counter -- silently bypassing the limit right around each +# rollover. Overwritten by each hop's own admission (last-write-wins), which +# is correct: success only ever fires for whichever hop actually served the +# request, so its own most recent admission timestamp is the right one. +_ADMISSION_TIME_FIELD: Final[str] = "_model_based_tag_rate_limits_admission_time" + +# The routing-group membership (candidate_model_names) admission actually +# resolved against, stashed the same way _ADMISSION_TIME_FIELD is. resolve_any +# dedupes divergent per-deployment entries by picking the alphabetically first +# member model_name sharing a signature (resolved_group) -- a pure function of +# this exact candidate set. Router's live routing-group membership can change +# between admission and success (a deployment added or removed mid-request via +# /model/new or a config hot-reload), and success independently re-deriving +# candidate_model_names from *live* membership at that later point can pick a +# different resolved_group than admission did, hashing to a different Redis +# key -- so success accounting silently misses the bucket admission actually +# checked, letting real usage escape the enforced cap. Reusing admission's own +# snapshot keeps resolve_any's output identical at both points regardless of +# what changed in between. +_ROUTING_GROUP_CANDIDATES_FIELD: Final[str] = "_model_based_tag_rate_limits_routing_group_candidates" + + +class _TagRateLimitIndex: + """Rebuilds the limits index when `llm_router.model_list` changes, or at + least every `_INDEX_TTL_SECONDS`, whichever comes first.""" + + def __init__(self, time_provider: Callable[[], datetime]) -> None: + self._time_provider = time_provider + self._cache_key: tuple[int, int] | None = None + self._built_at: float = 0.0 + self._index: _LimitsIndex = _LimitsIndex(by_model_name=MappingProxyType({}), by_team_alias=MappingProxyType({})) + + def get(self, llm_router: Router) -> _LimitsIndex: + model_list: Final = llm_router.model_list or () + cache_key: Final = (id(llm_router), len(model_list)) + now: Final = self._time_provider().timestamp() + if cache_key != self._cache_key or (now - self._built_at) >= _INDEX_TTL_SECONDS: + self._index = _build_limits_index(model_list) + self._cache_key = cache_key + self._built_at = now + return self._index + + +def _scope_suffix(deployment_scope: tuple[str, ...] | None) -> str: + return "chain" if deployment_scope is None else "dep:" + "+".join(deployment_scope) + + +def _hash_tag(model_group: str, configured: _ConfiguredLimit, tag_value: str, key_hash: str | None) -> str: + # resolved_group overrides the caller-visible model_group in two cases: + # resolve_any()'s per-deployment fallback (routing groups), where the + # caller-visible name is ambiguous (shared by every member model_name), so + # hashing by it would collide two different underlying model_names' + # identically-named limits onto one counter; and a team-owned deployment, + # where the caller-visible name differs depending on whether the caller + # dialed the team's public alias or the deployment's own internal + # model_name, so hashing by it would split one team's usage across two + # counters. See _ConfiguredLimit.resolved_group. + effective_model_group: Final = configured.resolved_group if configured.resolved_group is not None else model_group + scope: Final = _scope_suffix(configured.deployment_scope) + # team_scope disambiguates two teams that publish the identical + # team_public_model_name alias with identically-configured limits -- + # without it their buckets would collide despite the index correctly + # scoping the lookup by (team_id, alias). See _ConfiguredLimit.team_scope. + team_suffix: Final = f":team:{configured.team_scope}" if configured.team_scope is not None else "" + key_suffix: Final = f":key:{key_hash}" if key_hash is not None else "" + # Two entries can share `name`/`tag_id` while disagreeing on limit, + # period_seconds, or scoping (see _policy_fingerprint) -- included so + # they never collide onto the same counter despite the shared name. + policy_suffix: Final = f":policy:{_policy_fingerprint(configured.entry)}" + return ( + f"tag_rl:{effective_model_group}:{configured.unit}:{configured.entry.name}:{configured.entry.tag_id}:" + f"{scope}{team_suffix}:{_fixed_length_identity(tag_value)}{key_suffix}{policy_suffix}" + ) + + +def _bucket_key( + model_group: str, + configured: _ConfiguredLimit, + tag_value: str, + bucket_id: int, + key_hash: str | None = None, +) -> str: + return f"{{{_hash_tag(model_group, configured, tag_value, key_hash)}}}:{bucket_id}" + + +def _inflight_key( + model_group: str, + configured: _ConfiguredLimit, + tag_value: str, + key_hash: str | None = None, +) -> str: + """Concurrency counter key: not epoch-bucketed, since "how many are in + flight right now" has no window to reset on -- it's released explicitly + on completion, with a TTL fallback only for a leaked (crashed) reservation.""" + return f"{{{_hash_tag(model_group, configured, tag_value, key_hash)}}}:inflight" + + +class _ClassifiedCheck(NamedTuple): + configured_limit: _ConfiguredLimit + tag_value: str + key: str + is_atomic: bool + + +def _classify_check( + configured_limit: _ConfiguredLimit, + model: str, + tags: Sequence[str], + present_deployment_ids: frozenset[str], + request_kwargs: Mapping[str, object], + metadata_variable_name: str, + now: float, + key_alias: str | None, +) -> _ClassifiedCheck | None: + # A breach of a deployment-scoped check (this function's own + # `deployment_scope is not None` branch, and every one of its callers' + # `_raise_over_limit`) deliberately rejects the whole routing attempt for + # this hop, not just the deployment(s) that own it -- it does not filter + # them out of `healthy_deployments` and let a sibling in the same group + # serve instead. An earlier design considered filter-and-retry-sibling + # semantics (matching how native tag routing filters candidates rather + # than rejecting the hop) and deliberately did not adopt it: rejecting + # the whole hop is simpler, and avoids a caller silently succeeding + # against a deployment whose limit configuration they didn't intend to + # satisfy. + if configured_limit.deployment_scope is not None and not ( + present_deployment_ids & frozenset(configured_limit.deployment_scope) + ): + return None + tag_value: Final = _extract_identity(tags, configured_limit.entry.tag_id) + if tag_value is None: + return None + if not _entry_applies(configured_limit.entry, tags, key_alias, model): + return None + key_hash: Final = ( + _extract_key_hash(request_kwargs, metadata_variable_name) if configured_limit.entry.scope_by_key_hash else None + ) + if configured_limit.unit == "concurrency": + inflight_key: Final = _inflight_key(model, configured_limit, tag_value, key_hash=key_hash) + return _ClassifiedCheck(configured_limit, tag_value, inflight_key, is_atomic=True) + bucket_id: Final = int(now) // configured_limit.entry.period_seconds + bucket_key_value: Final = _bucket_key(model, configured_limit, tag_value, bucket_id, key_hash=key_hash) + return _ClassifiedCheck( + configured_limit, tag_value, bucket_key_value, is_atomic=configured_limit.unit in _ATOMIC_UNITS + ) + + +def _increment_operation_for_limit( + configured_limit: _ConfiguredLimit, + model_group: str, + tags: Sequence[str], + deployment_id: str | None, + key_hash: str | None, + key_alias: str | None, + increment_by_unit: Mapping[_LimitUnit, float], + now: float, +) -> RedisPipelineIncrementOperation | None: + if configured_limit.unit == "concurrency": + return None # released above, via _pop_pending_concurrency_keys + if configured_limit.deployment_scope is not None and deployment_id not in configured_limit.deployment_scope: + return None + tag_value: Final = _extract_identity(tags, configured_limit.entry.tag_id) + if tag_value is None: + return None + if not _entry_applies(configured_limit.entry, tags, key_alias, model_group): + return None + if configured_limit.unit not in increment_by_unit: + return None # "requests" is accounted atomically at admission, not here + increment_value: Final = increment_by_unit[configured_limit.unit] + if increment_value == 0: + return None + bucket_id: Final = int(now) // configured_limit.entry.period_seconds + key_hash_for_limit: Final = key_hash if configured_limit.entry.scope_by_key_hash else None + key: Final = _bucket_key(model_group, configured_limit, tag_value, bucket_id, key_hash=key_hash_for_limit) + return RedisPipelineIncrementOperation( + key=key, + increment_value=increment_value, + ttl=_bucket_ttl_seconds(configured_limit.entry), + ) + + +def _resolve_max_in_memory_cache_size() -> int | None: + """ + `litellm_settings` values reach `litellm.model_based_tag_rate_limits_max_in_memory_cache_size` + via a plain, unvalidated `setattr`, so a config typo (a negative number, or a + string like "500" from an unresolved os.environ/ substitution) can reach here. + InMemoryCache raises when comparing its size against a non-positive-int + max_size_in_memory, and DualCache.async_set_cache swallows that exception, so + an invalid value would otherwise silently disable every counter write for this + hook rather than fail loudly -- rejected here in favor of the safe default instead. + """ + configured: Final = litellm.model_based_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( + "model_based_tag_rate_limits_hook: model_based_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 + + +def _queue_pending_reservations( + request_kwargs: Mapping[str, object], field: str, reservations: Sequence[tuple[str, _PartitionKey]] +) -> None: + """Stash reservations on the request's own `model_call_details`, under + `field` -- see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring for why this, + not a `litellm_call_id`. Silently a no-op without a real logging object + (defensive only; every real request has one): a queued concurrency + reservation still self-heals via `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, + just later. + + Each entry is stamped with `_current_admission_token()`: `Router.abatch_completion`'s + comma-separated multi-model dispatch runs several `acompletion` calls + concurrently as *separate* tasks that all share one `model_call_details` + (the proxy attaches one `litellm_logging_obj` to the request before the + comma-split, and every branch inherits that same reference), so this + field is no longer scoped to one logical request's own serial fallback + chain the way its docstring assumes. `_release_stale_hop_reservations` + and the terminal release hooks use the stamp to tell "this same + admission lineage, safe to reclaim" apart from "a concurrent sibling + branch's own still-live reservation," which must never be touched here. + """ + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if not isinstance(model_call_details, dict): + return + pending = model_call_details.get(field) # rebind-ok: lazily initialized below when absent + if pending is None: + pending = [] # mutable-ok: shared, request-scoped accumulator; see field's own docstring # rebind-ok: lazily initialized only when absent + model_call_details[field] = pending + admission_token: Final = _current_admission_token() + pending.extend( + (key, partition_key, admission_token) for key, partition_key in reservations + ) # mutable-ok: see comment above + + +def _record_admission_time(request_kwargs: Mapping[str, object], now: float) -> None: + """Stash this hop's admission timestamp -- see `_ADMISSION_TIME_FIELD`'s + docstring for why. Silently a no-op without a real logging object + (defensive only; every real request has one): success accounting falls + back to its own fresh timestamp, same as before this fix existed.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if isinstance(model_call_details, dict): + model_call_details[_ADMISSION_TIME_FIELD] = now + + +def _admission_time_or(kwargs: Mapping[str, object], fallback: float) -> float: + recorded: Final = kwargs.get(_ADMISSION_TIME_FIELD) + return recorded if isinstance(recorded, float) else fallback + + +def _record_routing_group_candidates( + request_kwargs: Mapping[str, object], candidate_model_names: tuple[str, ...] +) -> None: + """Stash the routing-group membership admission resolved against -- see + `_ROUTING_GROUP_CANDIDATES_FIELD`'s docstring for why. Silently a no-op + without a real logging object (defensive only; every real request has + one): success accounting falls back to its own live reconstruction, same + as before this fix existed.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if isinstance(model_call_details, dict): + model_call_details[_ROUTING_GROUP_CANDIDATES_FIELD] = candidate_model_names + + +def _routing_group_candidates_or(kwargs: Mapping[str, object], fallback: tuple[str, ...]) -> tuple[str, ...]: + recorded: Final = kwargs.get(_ROUTING_GROUP_CANDIDATES_FIELD) + return recorded if isinstance(recorded, tuple) else fallback + + +@dataclass(frozen=True, slots=True) +class _CachePartition: + internal_usage_cache: InternalUsageCache + v3: _PROXY_MaxParallelRequestsHandler_v3 + + +class _PROXY_ModelBasedTagRateLimitsHook( # 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: + # A dedicated in-memory layer, not the proxy-wide `internal_usage_cache` + # passed in: that instance is shared with the key/team parallel-request + # limiter's own authentication-bound counters, and its default + # InMemoryCache evicts at 200 items. Without this isolation, a caller + # flooding this hook's own caller-controlled tag buckets past that + # ceiling could evict an unrelated, authentication-bound counter and + # exceed a limit nothing here configured. The real Redis connection + # (if any) is still shared across every partition (see _build_partition), + # so cross-instance correctness is unaffected regardless of partitioning. + self._redis_cache: Final = internal_usage_cache.redis_cache + self._time_provider = time_provider or datetime.now + # Every distinct _PartitionKey gets its own dedicated partition + # (in-memory cache + its own v3 handler), lazily built and memoized + # here -- see _partition_for. None (the key every entry uses unless + # it sets its own max_in_memory_cache_size) is this hook's single + # default partition, sized by + # litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200 if that's + # also unset), matching today's behavior for every entry that doesn't + # opt into its own partition. + 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._v3 = default_partition.v3 + self._index = _TagRateLimitIndex(time_provider=self._time_provider) + self._lock = asyncio.Lock() + self.llm_router: Router | None = None + 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 + ) + + def update_variables(self, llm_router: Router) -> None: + self.llm_router = llm_router + + 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, refresh_ttl: bool + ) -> tuple[bool, float]: + """Single-key atomic check-and-increment. Always one key per Lua + call -- see TAG_RL_CHECK_AND_INCR_SCRIPT's module docstring for why, + and for why `refresh_ttl` must be True for a concurrency key and + False for a requests key.""" + if self._check_and_incr_script is not None: + raw: Final = await self._check_and_incr_script( + keys=(key,), args=(limit, increment, ttl, 1 if refresh_ttl else 0) + ) + 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, refresh_ttl=refresh_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, bool]], + ) -> tuple[int | None, tuple[float, ...]]: + """ + All-or-nothing across every (cache, key, limit, increment, ttl) in + `checks`: if any would exceed its limit, none are incremented -- a + single hop's requests-unit and concurrency-unit checks must commit + together or not at all, even when they span more than one cache + partition. Each key is checked/incremented in its own single-key Lua + call (cluster-safe by construction); all-or-nothing across the batch + is enforced here by refunding every earlier admission the moment a + later key is rejected, not by a single multi-key script call. + + Refunds are best-effort: a refund that fails (e.g. a transient Redis + error) is logged and skipped rather than raised, so one bad refund + can't stop the rest of the batch from being refunded, and can't turn + a clean rejection into an unhandled exception. A skipped refund + self-heals via the key's TTL -- see `_ttl_for`. + + A later key's own admission raising (a transient Redis error, or + this coroutine being cancelled mid-call, e.g. the caller + disconnecting) is treated the same as a normal rejection for refund + purposes: only the earlier admissions in this batch are refunded, + never the raising key's own key. This is deliberate, not an + oversight: a raise gives no guarantee that key's own increment + didn't already commit server-side (Redis can run the INCRBY and + still have the call raise if the response back to us is lost), but + these are shared, chain-wide buckets with no per-request ownership + tracking -- decrementing on that guess is just as likely to erase a + *different*, legitimately-admitted concurrent request's charge on + the same key as it is to undo our own. That failure mode (an + attacker repeatedly cancelling requests to erase other callers' + charges and exceed the configured limit) is worse than the + alternative this accepts instead: a key that did commit but never + gets refunded self-heals via its own TTL -- see `_ttl_for`. The + earlier admissions refunded here are never ambiguous like this: they + are this same request's own confirmed-successful increments, so + undoing them is always safe. + + Refunds are best-effort: a refund that fails (e.g. a transient Redis + error) is logged and skipped rather than raised, so one bad refund + can't stop the rest of the batch from being refunded, and can't turn + a clean rejection into an unhandled exception. + + Returns (failing_index, values). On success, failing_index is None + and values holds each key's new post-increment value, same order as + `checks`. On rejection, failing_index is the 0-based index of the + first key that would have exceeded its limit and values holds that + one key's current (unmodified) value. + """ + if not checks: + return None, () + + # Sequential async admission: each element needs its own awaited + # Redis round trip, and a rejection mid-loop discards everything + # accumulated so far in favor of refunding and returning early, so + # this can't be expressed as a one-shot comprehension. + admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection; see comment above + for index, (cache, key, limit, increment, ttl, refresh_ttl) in enumerate(checks): + admitted = False + try: + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) + finally: + # Runs on a normal rejection (admitted stays False) and on + # any exception/cancellation from the awaited call above + # (admitted never gets assigned, so it's still the False set + # just before the try) -- either way, only the earlier, + # known-safe admissions are refunded; see the docstring + # above for why this key's own ambiguous outcome is not. + 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, bool]], up_to_index: int + ) -> None: + for refund_index in range(up_to_index): + refund_cache, refund_key, _limit, refund_increment, _ttl, _refresh_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( + "model_based_tag_rate_limits_hook: failed to refund %s on rollback: %s", refund_key, e + ) + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list, # mutable-ok: must match CustomLogger's base signature exactly, or basedpyright flags reportIncompatibleMethodOverride + messages: list[AllMessageValues] | None, # mutable-ok: see reason above + request_kwargs: dict | None = None, # mutable-ok: see reason above + parent_otel_span: Span | None = None, + ) -> list[dict]: # mutable-ok: see reason above + if ( + not healthy_deployments + or not isinstance(healthy_deployments, list) # pyright: ignore[reportUnnecessaryIsInstance] # defensive at runtime despite the static list annotation Router's own callers aren't guaranteed to honor + or self.llm_router is None + ): + return healthy_deployments + + resolved_request_kwargs: Final = request_kwargs or _EMPTY_MAPPING + stale_request_keys: Final = await self._release_stale_hop_reservations(resolved_request_kwargs) + # Not `get_metadata_variable_name_from_kwargs` (naive key-presence + # check): a caller can forge an empty (or `None`) `litellm_metadata` + # on an ordinary request to make that check pick it over the real, + # populated `metadata` the proxy wrote authenticated team/tag identity + # into, seeing no tags at all and admitting past every configured + # limit. See `_resolve_authoritative_metadata_variable_name`'s own + # docstring. + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(resolved_request_kwargs) + team_id: Final = _extract_team_id(resolved_request_kwargs, metadata_variable_name) + # Built from the full routing-group membership, not `healthy_deployments` + # (Router's own cooldown-filtered list for this hop): a member that's + # merely cooled down right now is still a real member of the group for + # the purpose of deciding resolved_group, and success accounting has no + # way to know which members were healthy at admission time -- it can + # only reconstruct the full, static membership (see its own comment + # below). Deriving both sides from the same full-membership source + # handles cooldown-state drift between the two points in time; actual + # membership drift (a deployment added or removed mid-request) still + # needs admission's own snapshot stashed and reused -- see + # _ROUTING_GROUP_CANDIDATES_FIELD's docstring. + routing_group_deployments: Final = self.llm_router._get_routing_group_deployments( # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching resolve_any's own reliance on this method + model=model, team_id=team_id + ) + candidate_model_names: Final = ( + tuple(dep["model_name"] for dep in routing_group_deployments) + if routing_group_deployments is not None + else tuple(name for d in healthy_deployments if isinstance(name := d.get("model_name"), str)) + ) + _record_routing_group_candidates(resolved_request_kwargs, candidate_model_names) + configured: Final = self._index.get(self.llm_router).resolve_any(model, team_id, candidate_model_names) + if not configured: + return healthy_deployments + + tags: Final = _order_tags_for_identity_resolution( + _get_tags_from_request_kwargs(resolved_request_kwargs, metadata_variable_name=metadata_variable_name), + resolved_request_kwargs, + metadata_variable_name, + ) + + present_deployment_ids: Final[frozenset[str]] = frozenset( + dep_id for d in healthy_deployments if (dep_id := _deployment_id(d)) is not None + ) + + key_alias: Final = _extract_key_alias(resolved_request_kwargs, metadata_variable_name) + now: Final = self._time_provider().timestamp() + _record_admission_time(resolved_request_kwargs, now) + classified: Final = tuple( + check + for configured_limit in configured + if ( + check := _classify_check( + configured_limit, + model, + tags, + present_deployment_ids, + resolved_request_kwargs, + metadata_variable_name, + now, + key_alias, + ) + ) + is not None + ) + read_only_checks: Final = tuple((c.configured_limit, c.tag_value, c.key) for c in classified if not c.is_atomic) + atomic_checks: Final = tuple((c.configured_limit, c.tag_value, c.key) for c in classified if c.is_atomic) + + current_values: Final = await self._read_only_values(read_only_checks, parent_otel_span) + 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 (a genexpr can't `await` here); zipped with atomic_checks immediately below + for configured_limit, _tag_value, _key in atomic_checks: + atomic_partitions_list.append( + await self._partition_for(_partition_key(configured_limit.entry)) + ) # mutable-ok: see comment above + atomic_partitions: Final = tuple(atomic_partitions_list) + failing_index, values = await self._atomic_check_and_increment( + tuple( + ( + partition.internal_usage_cache, + key, + configured_limit.entry.limit, + # A "requests" key matching one already charged by a + # superseded earlier hop of this same request (see + # _release_stale_hop_reservations) renews that same + # charge at zero net cost instead of adding a second + # unit on top of it -- folded into this same + # all-or-nothing batch so a hop that goes on to fail + # a *different* check here never commits a refund + # with nothing to replace it. + 0.0 if configured_limit.unit == "requests" and key in stale_request_keys else 1.0, + self._ttl_for(configured_limit), + configured_limit.unit == "concurrency", + ) + for partition, (configured_limit, _tag_value, key) in zip(atomic_partitions, atomic_checks) + ) + ) + if failing_index is not None: + failing_limit, failing_tag_value, _ = atomic_checks[failing_index] + self._raise_over_limit(failing_limit, failing_tag_value, model, current=values[0]) + + concurrency_reservations: Final = tuple( + (key, _partition_key(configured_limit.entry)) + for configured_limit, _tag_value, key in atomic_checks + if configured_limit.unit == "concurrency" + ) + if concurrency_reservations: + _queue_pending_reservations( + resolved_request_kwargs, _PENDING_CONCURRENCY_KEYS_FIELD, concurrency_reservations + ) + await self._mirror_pending_reservations( + resolved_request_kwargs.get("litellm_call_id"), + _extract_key_hash(resolved_request_kwargs, metadata_variable_name), + concurrency_reservations, + ) + + # Only genuinely new keys, never one already in stale_request_keys: + # that key's own check just renewed at zero net cost above and is + # still sitting in the field (see _release_stale_hop_reservations' + # own comment on why this is a peek, not a pop) -- appending it + # again here would grow the list with a duplicate entry on every + # hop of a long retry chain without changing what it means. + request_increments: Final = tuple( + (key, _partition_key(configured_limit.entry)) + for configured_limit, _tag_value, key in atomic_checks + if configured_limit.unit == "requests" and key not in stale_request_keys + ) + if request_increments: + _queue_pending_reservations( + resolved_request_kwargs, _PENDING_REQUEST_INCREMENTS_FIELD, request_increments + ) + + return healthy_deployments + + async def _mirror_pending_reservations( + self, call_id: object, key_hash: str | None, reservations: Sequence[tuple[str, "_PartitionKey"]] + ) -> None: + # Merged onto whatever's already mirrored, not overwritten: `call_id` + # and `key_hash` are identical across every branch of one + # abatch_completion dispatch (see _PENDING_CONCURRENCY_KEYS_FIELD's + # docstring), so an overwrite here would erase a still-live sibling + # branch's own mirrored reservation the moment this hop admits. + # Read-then-write is not atomic against a concurrent sibling doing + # the identical merge, so a rare interleaving can still lose one + # side's addition -- narrower than the guaranteed loss an overwrite + # caused, and self-heals via the same TTL either way. + if not isinstance(call_id, str): + return + cache_key: Final = _pending_reservations_cache_key(call_id, key_hash) + try: + existing: Final = _decode_reservations( + await self.internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None) + ) + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=_encode_reservations(existing + tuple(reservations)), + ttl=_CONCURRENCY_MIN_SAFETY_TTL_SECONDS, + litellm_parent_otel_span=None, + ) + except Exception as e: # noqa: BLE001 - a failed mirror write must never block admission; the reservation still self-heals via its own TTL + verbose_proxy_logger.warning( + "model_based_tag_rate_limits_hook: failed to mirror pending reservations for call_id=%s: %s", call_id, e + ) + + @staticmethod + def _ttl_for(configured_limit: _ConfiguredLimit) -> int: + if configured_limit.unit == "concurrency": + # A reservation's TTL must comfortably outlast any real in-flight + # request, or a slow request's reservation self-heals (expires) + # while it is still genuinely running, silently admitting extra + # requests past the configured limit. period_seconds (or an + # explicit key_ttl_seconds override) is still honored if the + # operator wants an even longer safety margin, but this floor is + # never lowered below it, even by an explicit override. + entry: Final = configured_limit.entry + 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(configured_limit.entry) + + async def _read_only_values( + self, + read_only_checks: Sequence[tuple[_ConfiguredLimit, str, str]], + parent_otel_span: Span | None, + ) -> tuple[float | None, ...]: + if not read_only_checks: + return () + + # Grouped by cache partition (one batched read per partition), then + # reassembled back into read_only_checks's original order: a hop can + # mix entries from more than one partition (e.g. a default-cache + # dollar_limits entry alongside a dedicated-partition request_limits + # entry), and _raise_if_over_limit below zips this result positionally + # against read_only_checks, so order must be preserved exactly. + indices_by_partition: Final[dict[_PartitionKey, list[int]]] = {} # mutable-ok: grouped, reassembled below + for index, (configured_limit, _tag_value, _key) in enumerate(read_only_checks): + partition_key = _partition_key(configured_limit.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(): + # not `Final`: rebound each loop iteration, which basedpyright's + # LIT010/Final-in-loop check forbids + partition = await self._partition_for(partition_key) + keys = [read_only_checks[i][2] 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: + # async_log_success_event increments these buckets straight + # through a Lua script on this same redis_cache, bypassing + # DualCache/InternalUsageCache entirely -- so its in-memory + # layer never learns about that write. DualCache's own + # async_batch_get_cache treats any non-None in-memory hit as + # authoritative and never re-checks Redis for that key (see + # _reserve_redis_batch_keys), so once a key is backfilled + # in-memory it silently freezes for up to the in-memory TTL + # (10 minutes by default) while the real Redis counter keeps + # moving underneath it -- reading straight from Redis here, + # bypassing that in-memory layer, is the only way this + # read-then-later-increment split stays coherent. + # not `Final`: rebound each loop iteration, which basedpyright's + # LIT010/Final-in-loop check forbids; explicitly typed (as the + # read-only supertype, since this is never mutated) since + # RedisCache.async_batch_get_cache's own signature returns a + # bare, unparameterized dict + 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[tuple[_ConfiguredLimit, str, str]], + current_values: Sequence[float | None], + model: str, + ) -> None: + for (configured_limit, tag_value, _key), current_value in zip(read_only_checks, current_values): + current = float(current_value) if current_value is not None else 0.0 + if current < configured_limit.entry.limit: + continue + self._raise_over_limit(configured_limit, tag_value, model, current=current) + + def _raise_over_limit( + self, + configured_limit: _ConfiguredLimit, + tag_value: str, + model: str, + current: float, + ) -> None: + verbose_proxy_logger.debug( + "model_based_tag_rate_limits_hook: OVER_LIMIT model=%s unit=%s name=%s tag_id=%s tag_value=%s current=%s limit=%s", + model, + configured_limit.unit, + configured_limit.entry.name, + configured_limit.entry.tag_id, + tag_value, + current, + configured_limit.entry.limit, + ) + raise ProxyRateLimitError( + detail={ # mutable-ok: must stay a real dict -- async_log_failure_event below (and generic proxy exception rendering, e.g. proxy/utils.py, guardrail hooks) branch on isinstance(exc.detail, dict); a MappingProxyType silently falls through those checks + "error": "tag_rate_limit_exceeded", + "type": configured_limit.unit, + "tag_id": configured_limit.entry.tag_id, + "tag_value": tag_value, + "limit_name": configured_limit.entry.name, + "limit": configured_limit.entry.limit, + "period_seconds": configured_limit.entry.period_seconds, + }, + headers={"retry-after": str(configured_limit.entry.period_seconds)}, # mutable-ok: same as detail + rate_limit_type=_UNIT_TO_RATE_LIMIT_TYPE[configured_limit.unit], + model=model, + llm_provider="litellm_proxy", + ) + + async def _release_keys(self, reservations: Sequence[tuple[str, _PartitionKey]]) -> None: + """ + Release each key by one slot. This does not verify the completing + request still owns a live reservation (no per-request slot identity + is tracked -- see the concurrency design note above), so a request + that outlives the safety TTL and gets its key reused by a fresh + reservation could in principle decrement a reservation it never + held. Flooring at 0 (TAG_RL_DECR_FLOOR_ZERO_SCRIPT) bounds the + damage to under-counting (briefly under-enforcing the limit) rather + than a negative counter, which would admit unlimited requests. + + Each reservation is released against the exact cache partition + (`_partition_for(partition_key)`) its increment used -- see + `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring for why this must match. + """ + 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( + "model_based_tag_rate_limits_hook: failed to release concurrency slot %s: %s", key, e + ) + + async def _release_stale_hop_reservations(self, request_kwargs: Mapping[str, object]) -> frozenset[str]: + """ + A concurrency reservation still queued when a *new* hop's admission + runs, *within the same admission lineage*, can only belong to an + earlier hop of this same request's own fallback chain that already + concluded and failed: Router awaits one hop's entire attempt (call + plus its own failure handling) before starting the next, and a hop + that instead succeeded ends the request there via + async_log_success_event, which already pops everything pending on + this same model_call_details -- so admission is never re-entered, in + that same lineage, while an earlier hop's reservation is still + legitimately in flight. + + The lineage check matters because `model_call_details` is not always + scoped to one such chain: `Router.abatch_completion`'s comma-separated + multi-model dispatch runs several branches concurrently, each its own + Task, but every branch shares the identical `litellm_logging_obj` (see + `_queue_pending_reservations`'s own docstring) -- so a reservation + queued by a still-running sibling branch can be sitting here too, and + releasing it out from under that branch would let more calls through + a concurrency limit than it allows. Only entries `_current_admission_token()` + stamped are safe to treat as stale; anything else is left for its own + branch to release. + + LiteLLM only invokes a request's CustomLogger.async_log_failure_event + once per request, for whichever hop fails first (its internal + has_logged_async_failure dedup silently skips every later hop's own + failure), so every hop after that one would otherwise never release + its predecessor's key until _CONCURRENCY_MIN_SAFETY_TTL_SECONDS. + Releasing here, at the one point guaranteed to re-run before every + subsequent hop, closes that gap for every hop except a final one + whose own failure exhausts the retry chain -- async_post_call_failure_hook + closes that residual case instead, via the cache mirror + `_PENDING_RESERVATIONS_CACHE_KEY_PREFIX` documents. + + A "requests" atomic increment queued under + `_PENDING_REQUEST_INCREMENTS_FIELD` is never refunded here, even + though the identical staleness invariant holds for it too: an + unconditional refund followed by this hop's own admission is not one + atomic operation, so a hop that goes on to fail a *different* check + (a read-only limit, or another entry in the same atomic batch) would + leave the refund committed with nothing to replace it, undercounting + a logical request that genuinely made an earlier, real attempt. The + returned keys let `async_filter_deployments` fold the swap into its + own atomic batch instead -- see its own comment for how. + + Deliberately a peek, not a pop, for that same field: an earlier + version popped it here and only re-queued on a fully successful + atomic batch, so a hop that failed *before* reaching that point (a + read-only check, or a different entry in its own batch) silently + dropped the bookkeeping -- the real counter was untouched (0.0 + renewals roll back to a no-op), but the *next* hop's own peek would + come back empty, no longer recognize the key as already charged, and + charge a fresh unit on top of the one still sitting in the real + counter. Peeking leaves the field exactly as it was for whichever + hop reads it next, regardless of how many hops in between fail + before ever reaching their own successful queuing step. + """ + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if not isinstance(model_call_details, dict): + return frozenset() + release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_own_lineage=True) + if release_keys: + await self._release_keys(release_keys) + pending_request_increments: Final = model_call_details.get(_PENDING_REQUEST_INCREMENTS_FIELD) + if not isinstance(pending_request_increments, list): + return frozenset() + admission_token: Final = _current_admission_token() + return frozenset(key for key, _partition_key, token in pending_request_increments if token is admission_token) + + async def _pop_pending_concurrency_keys( + self, kwargs: Mapping[str, object], *, only_own_lineage: bool = False + ) -> tuple[tuple[str, _PartitionKey], ...]: + # Snapshot then remove only those exact entries, never a blanket + # clear: a sibling branch sharing this same request's + # model_call_details can still be live and appending concurrently + # (see the field's own docstring), so wiping the whole list here + # would silently strand that branch's reservation instead of + # releasing it later. `only_own_lineage` additionally excludes any + # entry a *different* admission lineage queued -- see + # `_release_stale_hop_reservations`'s own docstring for why that + # distinction, not just presence, decides what's actually stale. + pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) + if not isinstance(pending, list) or not pending: + return () + snapshot: Final = ( + tuple(entry for entry in pending if entry[2] is _current_admission_token()) + if only_own_lineage + else tuple(pending) + ) + for entry in snapshot: + try: + pending.remove(entry) + except ValueError: + pass + released: Final = tuple(entry[:2] for entry in snapshot) + if released: + await self._discard_from_mirror(kwargs, released) + return released + + async def _discard_from_mirror( + self, kwargs: Mapping[str, object], released: Sequence[tuple[str, _PartitionKey]] + ) -> None: + # Removes only the entries this pop actually released, not a blanket + # delete of the whole (call_id, key_hash) mirror: that key is shared + # across every branch of one abatch_completion dispatch (see + # _PENDING_CONCURRENCY_KEYS_FIELD's docstring), so wiping it here + # would silently drop a still-live sibling branch's own entry before + # async_post_call_failure_hook ever gets to read it. + call_id: Final = kwargs.get("litellm_call_id") + if not isinstance(call_id, str): + return + # Not `get_metadata_variable_name_from_kwargs` (naive key-presence + # check): at this point `kwargs` is `model_call_details`, which + # carries `litellm_metadata` present-but-`None` alongside the real, + # populated `metadata` for a standard request -- see + # `_resolve_authoritative_metadata_variable_name`'s own docstring. + litellm_params_raw: Final = kwargs.get("litellm_params") + litellm_params_for_metadata: Final = litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) + key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) + cache_key: Final = _pending_reservations_cache_key(call_id, key_hash) + try: + mirrored: Final = list( # mutable-ok: local working copy, discarded after removing exactly `released`'s own occurrences below + _decode_reservations( + await self.internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None) + ) + ) + for entry in released: + try: + mirrored.remove(entry) + except ValueError: + pass + if mirrored: + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=_encode_reservations(tuple(mirrored)), + ttl=_CONCURRENCY_MIN_SAFETY_TTL_SECONDS, + litellm_parent_otel_span=None, + ) + else: + await self.internal_usage_cache.dual_cache.async_delete_cache(cache_key) + except Exception as e: # noqa: BLE001 - a failed mirror update must never block the real release above + verbose_proxy_logger.warning( + "model_based_tag_rate_limits_hook: failed to update mirrored reservations for call_id=%s: %s", + call_id, + e, + ) + + async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + """ + A client disconnecting before the first streamed chunk raises + CancelledError/GeneratorExit, which bypasses both async_log_success_event + and async_log_failure_event below -- the only two places a concurrency + reservation queued during admission is normally popped and released. + Without this, the reservation would sit held until _CONCURRENCY_MIN_SAFETY_TTL_SECONDS + expires, letting a caller who repeatedly opens and immediately drops + streaming requests exhaust their own tag's concurrency limit for free. + + Deliberately not `only_own_lineage=True`: `Router.abatch_completion` + returns every branch's response together rather than a single stream, + so its concurrent-sibling-branch race this hook otherwise guards + against (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring) cannot + co-occur with a mid-stream disconnect here. + """ + logging_obj: Final = request_data.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if not isinstance(model_call_details, dict): + return + release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details) + if release_keys: + await self._release_keys(release_keys) + + async def async_post_call_failure_hook( + self, + request_data: dict, # mutable-ok: must match CustomLogger.async_post_call_failure_hook's own base signature exactly + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + """ + litellm's Logging object sets has_logged_async_failure=True after + the first hop of a fallback chain fails, which blocks + async_log_failure_event for every later hop (see + fallback_event_handlers.py's own docstring) -- so a chain's own + final, chain-exhausting failure never reaches that callback at all, + and _release_stale_hop_reservations only cleans up a stale + reservation when a *next* hop's admission runs, which never happens + after the last one. This hook fires exactly once per proxy request, + at the point the proxy gives up and returns an error to the caller, + regardless of how many hops ran or whether the completion-level + callback was suppressed for this one. + + Reads the cache mirror written by `_mirror_pending_reservations`, not + `model_call_details`: proxy/utils.py's post_call_failure_hook pops + `litellm_logging_obj` off `request_data` before invoking any callback + here ("Remove before callbacks iterate — not serialisable"), and + neither a ContextVar nor `request_data` itself survives to this + point either (see `_PENDING_RESERVATIONS_CACHE_KEY_PREFIX`'s own + docstring for why, confirmed live for each). + + Keyed by `user_api_key_dict.api_key`, not a value read out of + `request_data`: the proxy's own auth middleware establishes + `user_api_key_dict` before any hook runs, so it can't be forged the + way `request_data["litellm_call_id"]` (the `x-litellm-call-id` + header) can -- see `_PENDING_RESERVATIONS_CACHE_KEY_PREFIX`'s + docstring for what a caller-forgeable-only key would let a caller do. + """ + call_id: Final = request_data.get("litellm_call_id") + if not isinstance(call_id, str): + return + cache_key: Final = _pending_reservations_cache_key(call_id, user_api_key_dict.api_key) + try: + raw: Final = await self.internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None) + except Exception as e: # noqa: BLE001 - a failed mirror read must never raise into the caller's request path + verbose_proxy_logger.warning( + "model_based_tag_rate_limits_hook: failed to read mirrored reservations for call_id=%s: %s", call_id, e + ) + return + release_keys: Final = _decode_reservations(raw) + if not release_keys: + return + try: + await self.internal_usage_cache.dual_cache.async_delete_cache(cache_key) + except Exception as e: # noqa: BLE001 - a failed mirror clear must never block the real release below + verbose_proxy_logger.warning( + "model_based_tag_rate_limits_hook: failed to clear mirrored reservations for call_id=%s: %s", call_id, e + ) + 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: a hop whose own admission rejects never reaches the + # point where a concurrency reservation is queued (see + # async_filter_deployments), so _pop_pending_concurrency_keys already + # returns nothing to release in that case. Skipping release based on + # the exception's error marker alone would be wrong here, since + # global_tag_rate_limits_hook raises the identical marker -- that + # rejection can land after this hook already reserved a slot for the + # same request, and that slot must still be released. + # + # Not `only_own_lineage=True`: a streaming response is consumed (and + # this event fired) from a task the proxy forks independently of + # admission's own, so `_ADMISSION_CONTEXT` never reaches it either -- + # requiring a match here would leak every streaming request's + # reservation until `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS` instead of + # releasing it. `async_release_disconnect_state_hook` stays + # unfiltered for the identical reason; releasing everything here + # keeps a still-live `abatch_completion` sibling's own reservation + # exposed to the same risk that hook already accepts. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) + if release_keys: + await self._release_keys(release_keys) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + # Not `only_own_lineage=True`: see async_log_failure_event's own comment above. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) + if release_keys: + 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) + + if self.llm_router is None: + return + + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") + if standard_logging_object is None: + return + + model_group: Final = standard_logging_object.get("model_group") + if not model_group: + 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_authoritative_metadata_variable_name(litellm_params_for_metadata) + team_id: Final = _extract_team_id(litellm_params_for_metadata, metadata_variable_name) + 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) + # model_group is the caller-visible name, which Router deliberately + # keeps distinct from the serving deployment's own model_name for a + # routing-group call (see resolve_any's docstring). Passing only the + # one deployment that actually served this hop as the sole candidate + # would make resolve_any's dedup independently re-derive a + # *different* resolved_group than admission did whenever the group + # has more than one member: admission sees every member and picks + # whichever one frozenset(candidate_model_names) yields first for a + # shared signature, so success accounting must reconstruct that same + # full candidate set to land on the identical bucket, not just + # whichever deployment happened to serve -- otherwise a token/dollar + # limit is checked against one bucket at admission and accounted + # against a different one on success, letting usage silently bypass + # the configured limit. Falls back to the serving deployment alone + # only when `model_group` isn't a routing group at all (a plain + # single-model_name chain, where resolve() already matches directly + # and this candidate set is never actually consulted). + # + # Reconstructing live membership here is itself only a fallback: it + # can still disagree with admission's own candidate set if the + # routing group's actual membership changed between the two points + # in time (not just cooldown/health state) -- _routing_group_candidates_or + # below prefers admission's own stashed snapshot whenever one exists. + # See _ROUTING_GROUP_CANDIDATES_FIELD's docstring. + deployment_id: Final = standard_logging_object.get("model_id") + serving_deployment: Final = ( + self.llm_router.get_deployment(deployment_id) if isinstance(deployment_id, str) else None + ) + routing_group_deployments: Final = self.llm_router._get_routing_group_deployments( # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching resolve_any's own reliance on this method + model=model_group, team_id=team_id + ) + live_candidate_model_names: Final = ( + tuple(dep["model_name"] for dep in routing_group_deployments) + if routing_group_deployments is not None + else ((serving_deployment.model_name,) if serving_deployment is not None else ()) + ) + candidate_model_names: Final = _routing_group_candidates_or(kwargs, fallback=live_candidate_model_names) + configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names) + if not configured: + return + + 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 + + now: Final = _admission_time_or(kwargs, fallback=self._time_provider().timestamp()) + 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_limit: Final = tuple( + (configured_limit, operation) + for configured_limit in configured + if ( + operation := _increment_operation_for_limit( + configured_limit, model_group, tags, deployment_id, key_hash, key_alias, increment_by_unit, now + ) + ) + is not None + ) + + if not operation_by_limit: + return + + # Grouped by cache partition: a hop's tokens/dollars entries can span + # more than one partition, and each partition owns its own v3 + # handler (see _build_partition), so each group's operations are + # pipelined through that partition's own handler. + operations_by_partition: Final[_PartitionOperations] = {} # mutable-ok: see comment above + for configured_limit, operation in operation_by_limit: + partition_key = _partition_key(configured_limit.entry) + operations = operations_by_partition.setdefault(partition_key, []) # mutable-ok: see above + operations.append(operation) # mutable-ok: see comment above + + parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) + 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=group_operations, + parent_otel_span=parent_otel_span, + ) + ) + _BACKGROUND_TASKS.add(accounting_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + accounting_task.add_done_callback(_BACKGROUND_TASKS.discard) diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py new file mode 100644 index 00000000000..dafc1b0a95d --- /dev/null +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -0,0 +1,401 @@ +""" +Primitives shared by both tag-scoped rate-limit hooks: +`model_based_tag_rate_limits_hook.py` (per-deployment limits nested under +`model_info.tag_rate_limits`, admitted once per routing hop) and +`global_tag_rate_limits_hook.py` (model-independent limits declared once +under `litellm_settings.global_tag_rate_limits`, admitted once per request). + +Both hooks enforce the identical `TagRateLimitEntry` shape and need the same +identity/scope extraction, policy fingerprinting, bucket-key hashing, and +cache-partitioning primitives, so those live here rather than in either +hook's own module -- keeping neither hook reaching into the other's private +internals to reuse them. Every name here is this module's own public +interface (no leading underscore): each hook imports what it needs aliased +back to its own historical, underscore-prefixed local name (e.g. +`entry_applies as _entry_applies`), so this is a genuine export, not a +private symbol either hook reaches across a module boundary to grab. +""" + +import asyncio +import hashlib +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from litellm.exceptions import RateLimitType +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import TagRateLimitEntry, TagRateLimitScope + +LimitUnit: TypeAlias = Literal["tokens", "requests", "dollars", "concurrency"] +LIMIT_UNITS: Final[tuple[LimitUnit, ...]] = ("tokens", "requests", "dollars", "concurrency") + +# Units whose admission must be atomic (check-and-increment in one Redis +# round trip) because the increment amount is known upfront (always 1). +# tokens/dollars can't be: real usage is only known after the response, so +# they stay a read-then-account-on-success check with a documented, +# unavoidable admit-vs-account race. +ATOMIC_UNITS: Final[frozenset[LimitUnit]] = frozenset({"requests", "concurrency"}) + +UNIT_TO_GROUP_FIELD: Final[Mapping[LimitUnit, str]] = MappingProxyType( + { + "tokens": "token_limits", + "requests": "request_limits", + "dollars": "dollar_limits", + "concurrency": "concurrency_limits", + } +) +UNIT_TO_RATE_LIMIT_TYPE: Final[Mapping[LimitUnit, RateLimitType]] = MappingProxyType( + { + "tokens": RateLimitType.TOKENS, + "requests": RateLimitType.REQUESTS, + "dollars": RateLimitType.BUDGET, + "concurrency": RateLimitType.CONCURRENT_REQUESTS, + } +) + +# Shared read-only fallback for an absent/None mapping (request_kwargs, +# metadata, model_info, ...): avoids constructing a fresh mutable `{}` at +# every one of these call sites just to immediately call `.get()` on it. +EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) + +# `asyncio.create_task`'s own docs: "Save a reference to the result of this +# function, to avoid a task disappearing mid-execution. The event loop only +# keeps weak references to tasks. A task that isn't referenced elsewhere may +# get garbage collected at any time, even before it's done." The success path +# deliberately fires-and-forgets its concurrency release and its token/dollar +# accounting increment (unlike the failure/disconnect paths, which await +# concurrency release directly) to keep the hot success-response path from +# waiting on a Redis round trip; by the time either background task would +# run, the state it needs (popped pending keys, or the request's own usage +# figures) is only available in that task's own closure, so a collected +# task's work is unrecoverable, not just delayed. Holding a strong reference +# here until each task's own completion callback discards it is the standard +# fix, shared by every fire-and-forget task either hook creates. +BACKGROUND_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: see comment above + +# Floor for a concurrency reservation's self-heal TTL, regardless of the +# configured period_seconds. A reservation that expires while its request is +# still genuinely in flight silently admits requests past the limit; this +# generous floor keeps that window far larger than any realistic request +# duration, at the cost of a leaked (crashed-worker) slot self-healing more +# slowly. period_seconds can still raise the TTL further, never lower it. +CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600 + +# Single-key atomic check-and-increment. Deliberately one key per script call +# (never a batch of differently-hash-tagged keys in one call): every tag_rl +# key carries its own self-contained {..} hash tag so unrelated buckets never +# forcibly co-locate on the same Redis Cluster shard, which means a single Lua +# invocation can never span more than one key's slot without risking a +# cross-slot error. All-or-nothing across a hop's multiple atomic checks +# (e.g. requests + concurrency checked together) is achieved in Python by +# calling this once per key and refunding every earlier admission in the same +# batch if a later one is rejected -- the same refund-on-rollback shape as +# `atomic_check_and_increment_by_n` in parallel_request_limiter_v3.py, applied +# per-key instead of per-descriptor since each key already is one hash-tag +# group by construction. +# +# refresh_ttl (ARGV[4]) distinguishes the two callers of this script: +# "requests" is an epoch-bucketed fixed window, whose TTL must be set once +# (at first write) and never extended, or the bucket outlives the epoch it's +# meant to reset at. "concurrency" is not windowed at all -- its TTL exists +# purely as a crash-safety net for a reservation whose explicit release never +# runs -- so a still-active bucket must keep pushing that TTL out on every +# admission, or a long-lived burst of continuous traffic expires the whole +# counter mid-flight (silently admitting past the cap, and letting a release +# for a since-reset counter decrement an unrelated, newer cohort). +TAG_RL_CHECK_AND_INCR_SCRIPT: Final = """ +local key = KEYS[1] +local limit = tonumber(ARGV[1]) +local increment = tonumber(ARGV[2]) +local ttl = tonumber(ARGV[3]) +local refresh_ttl = tonumber(ARGV[4]) +local current = tonumber(redis.call('GET', key) or 0) +if current + increment > limit then + return { 0, current } +end +local new_value = redis.call('INCRBY', key, increment) +if ttl > 0 then + if refresh_ttl == 1 then + redis.call('EXPIRE', key, ttl) + else + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl) + end + end +end +return { 1, new_value } +""" + +# Atomic decrement that never leaves a counter negative. Used both to refund +# an earlier admission when a later key in the same batch is rejected, and to +# release a concurrency reservation -- floors at 0 so a decrement that can't +# be attributed to the exact reservation that caused it (see +# `_release_keys`'s docstring) degrades to under-counting rather than a +# negative counter that would admit unlimited requests. Floors via DEL, not +# `SET key 0`: releasing a reservation whose key already expired makes +# INCRBY recreate it with no TTL, and a plain SET would leave that recreated +# key permanently in Redis (SET clears any TTL); DEL removes it outright, +# which reads back identically to 0 everywhere this key is read (`GET key or 0`). +TAG_RL_DECR_FLOOR_ZERO_SCRIPT: Final = """ +local key = KEYS[1] +local delta = tonumber(ARGV[1]) +local new_value = redis.call('INCRBY', key, delta) +if new_value < 0 then + redis.call('DEL', key) + new_value = 0 +end +return new_value +""" + +# A (tag_id, values) pair mirroring TagRateLimitScope's own fields, used to +# fold `enabled_for`/`disabled_for` into a hashable form (a policy +# fingerprint, or a per-deployment dedup signature) without depending on +# TagRateLimitScope's own hashability. +ScopeSignature: TypeAlias = tuple[str, tuple[str, ...]] | None + + +def scope_signature(scope: TagRateLimitScope | None) -> ScopeSignature: + """Normalizes a `TagRateLimitScope` into a plain, hashable tuple, so + `enabled_for`/`disabled_for` can be folded into a policy fingerprint (or + a per-deployment dedup signature) -- two entries disagreeing on either + field must be treated as genuinely different policies rather than + merged into one bucket.""" + return None if scope is None else (scope.tag_id, scope.values) + + +def extract_identity(tags: Sequence[str], tag_id: str) -> str | None: + """ + First tag matching `f"{tag_id}:"`, value after the colon. Tags starting + with `!` are tag-routing negation markers, not identity tags, and are + skipped so they can never be misread as an identity value. + """ + prefix: Final = f"{tag_id}:" + for tag in tags: + if tag.startswith("!"): + continue + if tag.startswith(prefix): + return tag[len(prefix) :] + return None + + +def entry_applies(entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, model: str | None) -> bool: + """ + Applies `entry`'s own scoping fields (`enabled_for`/`disabled_for`/ + `apply_to_key_alias`/`apply_to_models`), evaluated in this order -- deny + overrides allow, checked before any allowlist: + + 1. `disabled_for`: the gate tag (often a SECOND, independent tag, but + `disabled_for.tag_id` can equally be set to this entry's own + `tag_id` to gate on a subset of its own resolved identity) is + present and its value is in `disabled_for.values` -> doesn't apply. + Absent gate tag never triggers this -- nothing to match against a + denylist. + 2. `enabled_for`: the gate tag is absent, or present but its value is + NOT in `enabled_for.values` -> doesn't apply. Unlike `disabled_for`, + absence DOES fail this check -- an allowlist gate requires an + explicit match, so "not tagged at all" means "not in scope". + 3. `apply_to_models`: `model` is absent, or present but not in the + list -> doesn't apply. Same allowlist semantics as `enabled_for` -- + a request with no `model` never satisfies this gate. + 4. `apply_to_key_alias`: the calling key's own alias is absent, or + present but not in the list -> doesn't apply. Same allowlist + semantics as `enabled_for` -- a key with no alias set never + satisfies this gate. + + An entry with none of these fields set always applies -- this is the + unscoped behavior every existing entry has today, unchanged. + """ + if entry.disabled_for is not None: + disabled_gate_value: Final = extract_identity(tags, entry.disabled_for.tag_id) + if disabled_gate_value is not None and disabled_gate_value in entry.disabled_for.values: + return False + if entry.enabled_for is not None: + enabled_gate_value: Final = extract_identity(tags, entry.enabled_for.tag_id) + if enabled_gate_value is None or enabled_gate_value not in entry.enabled_for.values: + return False + if entry.apply_to_models is not None and model not in entry.apply_to_models: + return False + if entry.apply_to_key_alias is None: + return True + return key_alias in entry.apply_to_key_alias + + +def resolve_authoritative_metadata_variable_name( + metadata_source: Mapping[str, object], +) -> Literal["metadata", "litellm_metadata"]: + """`get_metadata_variable_name_from_kwargs` only checks key presence, which + misresolves both at admission time and at `async_log_success_event` time: + a caller can forge an empty (or merely present-but-`None`) `litellm_metadata` + on an ordinary request -- `kwargs["litellm_params"]` also always carries a + `litellm_metadata` key (typically `None`) alongside the real, populated + `metadata` dict for a standard (non LITELLM_METADATA_ROUTES) request -- and + the key-presence check always picks `litellm_metadata` in both cases, + silently reading no tags/identity at all and admitting the request against + every configured limit. + + Merely requiring the value to be a non-empty dict is not enough either: a + caller can populate its own, unrelated keys on the non-authoritative + bucket (e.g. `{"litellm_metadata": {"x": 1}}` on an ordinary route), which + is non-empty but still not the field the proxy wrote identity into. + `add_litellm_data_to_request` unconditionally stamps `user_api_key_auth` + into whichever bucket the route actually resolved as authoritative, and + strips any `user_api_key_`-prefixed key a caller pre-populates on the + other bucket -- so requiring that marker's presence, not mere + truthiness, can't be forged onto the wrong side.""" + litellm_metadata: Final = metadata_source.get("litellm_metadata") + if isinstance(litellm_metadata, Mapping) and "user_api_key_auth" in litellm_metadata: + return "litellm_metadata" + return "metadata" + + +def extract_key_hash(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Same single-authoritative-field lookup as + `model_based_tag_rate_limits_hook._extract_team_id`, but for the calling + virtual key's hash: `LiteLLMProxyRequestSetup` sets + `metadata["user_api_key"]` to `user_api_key_dict.api_key`, which despite + the plain name is already the hashed token (see `litellm_pre_call_utils.py`). + """ + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + key_hash: Final = active.get("user_api_key") + return key_hash if isinstance(key_hash, str) else None + + +def extract_key_alias(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Same single-authoritative-field lookup as + `model_based_tag_rate_limits_hook._extract_team_id`, but for the calling + virtual key's own `key_alias`: `LiteLLMProxyRequestSetup` sets + `metadata["user_api_key_alias"]` to `user_api_key_dict.key_alias` + (see `litellm_pre_call_utils.py`).""" + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + key_alias: Final = active.get("user_api_key_alias") + return key_alias if isinstance(key_alias, str) else None + + +def _active_metadata_bucket(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> Mapping[str, object]: + """Same fallback `_get_tags_from_request_kwargs` (tag_based_routing.py) + already relies on: `request_kwargs` is a flat, top-level-metadata dict at + admission time, but `Logging.model_call_details` (what `kwargs` actually + is by `async_log_success_event`/`async_log_failure_event` time) never + carries `metadata`/`litellm_metadata` at its own top level, only nested + under `request_kwargs["litellm_params"]`. Checking only the top level + silently finds nothing at success/failure time, exactly the same failure + mode that lookup already had to handle.""" + top_level: Final = request_kwargs.get(metadata_variable_name) + if isinstance(top_level, Mapping): + return top_level + litellm_params: Final = request_kwargs.get("litellm_params") + if isinstance(litellm_params, Mapping): + nested: Final = litellm_params.get(metadata_variable_name) + if isinstance(nested, Mapping): + return nested + return EMPTY_MAPPING + + +def order_tags_for_identity_resolution( + tags: Sequence[str], request_kwargs: Mapping[str, object], metadata_variable_name: str +) -> tuple[str, ...]: + """`extract_identity`/`entry_applies` both resolve a `tag_id` via + first-match-by-prefix. `_merge_tags` (litellm_pre_call_utils.py) appends + key/team/project tags only if not already present, keeping caller-supplied + tags first in the merged `tags` list -- so an authenticated caller could + submit e.g. `company_id:attacker-chosen` ahead of the calling key's real + `company_id:real-company` tag and have every entry scoped to `company_id` + resolve to the caller's own value instead of the key's. `metadata.inherited_tags` + is a separate, server-computed snapshot of only the tags the calling + key/team/project's own config contributed (see that field's docstring in + litellm_pre_call_utils.py), so putting it first makes a policy-backed tag + win over a same-prefix caller-supplied one. + """ + active: Final = _active_metadata_bucket(request_kwargs, metadata_variable_name) + inherited_tags: Final = active.get("inherited_tags") + if not isinstance(inherited_tags, (list, tuple)) or not inherited_tags: + return tuple(tags) + return tuple(dict.fromkeys((*inherited_tags, *tags))) + + +def fixed_length_identity(tag_value: str) -> str: + """ + `tag_value` is caller-controlled (whatever follows the tag_id prefix in + a caller-supplied tag) with no length or content bound. Embedding it + directly would let a caller inflate a hook's own in-memory dict keys + past what `max_in_memory_cache_size` bounds (that caps item *count*, not + key bytes) and grow unbounded Redis keys with no cap at all. Hashing to + a fixed-length digest bounds a hook's own contribution to key size + regardless of the caller's input, while still preserving distinctness + (two different tag values still resolve to two different buckets). + """ + return hashlib.sha256(tag_value.encode()).hexdigest() + + +def policy_fingerprint(entry: TagRateLimitEntry) -> str: + """ + Two entries can share a `name` and `tag_id` while genuinely disagreeing + on `limit`, `period_seconds`, `scope_by_key_hash`, or any of the scoping + fields -- both hooks already treat that as two distinct policies + elsewhere (model_based_tag_rate_limits_hook.py's own per-deployment + dedup signature is one example), so the Redis/in-memory bucket key must + too, or two differently-configured entries that happen to share a name + check and charge the identical counter. `scope_by_key_hash` specifically + needs its own slot here rather than relying on a hook's own `_hash_tag` + `key_hash`-derived suffix to carry it: that suffix is empty whenever + `key_hash` resolves to `None` (no virtual key on the call), which would + otherwise collide an unscoped entry with a key-hash-scoped one that + agrees on every other field. Hashed to a fixed-length digest for the + same reason `fixed_length_identity` hashes `tag_value`: an operator's + own `enabled_for`/`disabled_for`/`apply_to_key_alias`/`apply_to_models` + list has no length bound. + """ + fingerprint_source: Final = ( + entry.limit, + entry.period_seconds, + entry.scope_by_key_hash, + scope_signature(entry.enabled_for), + scope_signature(entry.disabled_for), + entry.apply_to_key_alias, + entry.apply_to_models, + ) + return hashlib.sha256(repr(fingerprint_source).encode()).hexdigest()[:16] + + +def bucket_ttl_seconds(entry: TagRateLimitEntry) -> int: + """Redis (and in-memory fallback) TTL for a non-concurrency bucket key. + `entry.key_ttl_seconds` overrides the default of period_seconds + 3600 + when set -- see TagRateLimitEntry.key_ttl_seconds.""" + return entry.key_ttl_seconds if entry.key_ttl_seconds is not None else entry.period_seconds + 3600 + + +# None => this entry shares the hook's single default cache partition +# (matching every entry's behavior before this override existed). Otherwise +# a value-stable signature -- not the override int alone -- so two different +# entries that happen to choose the identical max_in_memory_cache_size don't +# get merged into one shared partition; the same entry (same config content) +# always resolves to the same signature across config/index rebuilds, which +# is what keeps each hook's own `_partitions` cache from leaking a fresh +# partition every time its configuration is re-resolved. +# The str before the trailing int is `policy_fingerprint(entry)`: two +# entries can share tag_id/name while genuinely disagreeing on +# limit/period_seconds/scope_by_key_hash/enabled_for/disabled_for/ +# apply_to_key_alias/apply_to_models -- both hooks already treat that as two +# distinct policies for bucket-key purposes, so a shared +# max_in_memory_cache_size must not route them onto the same partition +# either, or one entry's high-cardinality traffic can evict the other's +# active counters from a cache neither entry asked to share. +# max_in_memory_cache_size stays the trailing element: `partition_key[-1]` +# reads it directly to size the partition's cache. +PartitionKey: TypeAlias = tuple[str, str, str, int] | None +# Grouping type for async_log_success_event's per-partition tokens/dollars +# pipeline dispatch -- named only so the declaration fits on one line; see +# that method for why the grouping is needed. +PartitionOperations: TypeAlias = dict[PartitionKey, list[RedisPipelineIncrementOperation]] + + +def partition_key(entry: TagRateLimitEntry) -> PartitionKey: + if entry.max_in_memory_cache_size is None: + return None + return ( + entry.tag_id, + entry.name, + policy_fingerprint(entry), + entry.max_in_memory_cache_size, + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index a3335be2b2b..47eae2bf4a8 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum +import math from collections.abc import Mapping from dataclasses import dataclass from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -137,6 +138,176 @@ def _as_utc(value: datetime.datetime | None) -> datetime.datetime | None: return value.astimezone(datetime.timezone.utc) +class TagRateLimitScope(BaseModel): + """ + A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + entry to `tag_id: company_id, values: ["1032"]` so it only applies to + requests tagged as belonging to company 1032, independent of whichever + tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + `disabled_for`, which are the only two fields that construct this. + """ + + tag_id: str + values: tuple[str, ...] + + model_config = ConfigDict(frozen=True) + + @model_validator(mode="after") + def _validate_values(self) -> "TagRateLimitScope": + if not self.values: + raise ValueError("values must be a non-empty list of strings") + return self + + @model_validator(mode="after") + def _normalize_values(self) -> "TagRateLimitScope": + # Sorted and deduplicated: `values` is only ever used for membership + # tests (see _entry_applies), never order-dependent, but is also + # folded verbatim into the dedup signature two deployments' entries + # are compared by (see _scope_signature) -- an unsorted tuple would + # make config-order alone, not policy, decide whether two entries + # dedup to one shared bucket or wrongly split into two. + # object.__setattr__ bypasses this frozen model's own assignment + # guard -- returning a replacement instance from an "after" validator + # is silently ignored when constructing via __init__ (only takes + # effect via model_validate), so mutating in place is the only way + # this normalization reliably applies regardless of construction path. + object.__setattr__(self, "values", tuple(sorted(set(self.values)))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitEntry(BaseModel): + name: str + tag_id: str = "end_user_id" + limit: float + period_seconds: int + scope_by_key_hash: bool = False + # Overrides this entry's bucket/reservation key TTL (Redis, and the + # in-memory fallback when Redis isn't configured). Defaults to + # period_seconds + 3600 when unset -- see _PROXY_ModelBasedTagRateLimitsHook._ttl_for. + # A high-cardinality tag_id can keep many keys alive at once; lowering + # this lets an operator shed them sooner without shortening + # period_seconds itself. + key_ttl_seconds: int | None = None + # Overrides the size of the dedicated in-memory cache partition this + # entry's own keys live in, when Redis isn't configured (or as a local + # fast-path cache when it is). Unset means this entry shares the hook's + # single default partition, sized by + # litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200 if that's also + # unset). A high-cardinality tag_id can churn past that shared cap and + # evict another entry's active counters; setting this gives the entry + # its own dedicated partition instead. + max_in_memory_cache_size: int | None = None + # Gate this entry on a tag -- often a SECOND, independent tag (e.g. + # `enabled_for: {tag_id: company_id, values: ["1032"]}` to scope an + # override to one company's traffic), but `tag_id` can equally be set to + # this same entry's own `tag_id` to scope by a subset of its own + # resolved identity instead, without a second tag at all. + # `disabled_for` is checked first (deny overrides allow) when both are + # set. An absent gate tag never satisfies `enabled_for` (an allowlist + # gate requires an explicit match) but never triggers `disabled_for` + # either (nothing to match against a denylist). + enabled_for: TagRateLimitScope | None = None + disabled_for: TagRateLimitScope | None = None + # Restrict this entry to requests authenticated with one of these virtual + # keys' own `key_alias`. Unset (the default) means the entry applies to + # every request regardless of which key made it. A key with no alias set + # never satisfies this allowlist, same "absent gate never matches an + # allowlist" precedent as `enabled_for`. + apply_to_key_alias: tuple[str, ...] | None = None + # Restrict this entry to requests whose caller-facing `model` matches one + # of these names. Unset (the default) means the entry applies to every + # model. A request with no `model` field never satisfies this allowlist, + # same "absent gate never matches an allowlist" precedent as + # `apply_to_key_alias`. + apply_to_models: tuple[str, ...] | None = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _validate_limit(self) -> "TagRateLimitEntry": + # NaN compares False against every ordering operator, so a NaN limit + # makes the atomic requests/concurrency check-and-increment (which + # rejects when the new value exceeds the limit) never reject -- + # admitting indefinitely -- while the read-only tokens/dollars check + # (which admits when the current value is under the limit) never + # admits, rejecting every tagged request. Either outcome silently + # defeats the entry; reject it at config load time instead. + if math.isnan(self.limit): + raise ValueError("limit must not be NaN") + # Positive infinity never rejects the checks that gate this limit; negative + # infinity always does. Both silently defeat the entry. + if math.isinf(self.limit): + raise ValueError("limit must be finite") + # Zero or negative makes every check that gates this limit either always + # reject or never admit, silently blocking or admitting all matching traffic + # instead of the likely intended config. + if self.limit <= 0: + raise ValueError("limit must be a positive number") + return self + + @model_validator(mode="after") + def _validate_period_seconds(self) -> "TagRateLimitEntry": + if self.period_seconds <= 0: + raise ValueError("period_seconds must be a positive integer") + return self + + @model_validator(mode="after") + def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": + if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: + raise ValueError("key_ttl_seconds must be a positive integer when set") + # A shorter TTL than period_seconds expires the counter before its period + # elapses, letting it reset early and exceed the limit. + if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: + raise ValueError("key_ttl_seconds must be at least period_seconds when set") + return self + + @model_validator(mode="after") + def _validate_max_in_memory_cache_size(self) -> "TagRateLimitEntry": + if self.max_in_memory_cache_size is not None and self.max_in_memory_cache_size <= 0: + raise ValueError("max_in_memory_cache_size must be a positive integer when set") + return self + + @model_validator(mode="after") + def _validate_apply_to_key_alias(self) -> "TagRateLimitEntry": + if self.apply_to_key_alias is not None and not self.apply_to_key_alias: + raise ValueError("apply_to_key_alias must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_key_alias(self) -> "TagRateLimitEntry": + # Sorted and deduplicated for the same reason as + # TagRateLimitScope._normalize_values: only ever used for membership + # tests, but also folded verbatim into the dedup signature, where an + # unsorted tuple would make config-order alone decide whether two + # deployments' entries dedup to one shared bucket. + if self.apply_to_key_alias is not None: + self.apply_to_key_alias = tuple(sorted(set(self.apply_to_key_alias))) # mutable-ok: frozen before escaping + return self + + @model_validator(mode="after") + def _validate_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None and not self.apply_to_models: + raise ValueError("apply_to_models must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None: + self.apply_to_models = tuple(sorted(set(self.apply_to_models))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitGroup(BaseModel): + limits: tuple[TagRateLimitEntry, ...] = () + + +class TagRateLimits(BaseModel): + token_limits: TagRateLimitGroup | None = None + request_limits: TagRateLimitGroup | None = None + dollar_limits: TagRateLimitGroup | None = None + concurrency_limits: TagRateLimitGroup | None = None + + class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. @@ -186,6 +357,8 @@ class ModelInfo(MirroredPricingParams): # router-wide default. enable_tag_filtering: bool | None = None + tag_rate_limits: TagRateLimits | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 85e8308ae91..1e03159ca9a 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -45,6 +45,39 @@ async def test_async_increment_delegates_to_locked_sync_path(): assert cache.get_cache("counter") == 5 +async def test_async_increment_pipeline_applies_each_operations_own_ttl(): + """ + Bugbot finding: async_increment_pipeline dropped each operation's own + "ttl" field, so a counter created through it always got the cache's + 600-second default_ttl regardless of what the caller actually configured + (e.g. an hourly or daily rate-limit window) -- the counter (and whatever + it was tracking against a limit) silently reset mid-window. + """ + cache = InMemoryCache(default_ttl=600) + await cache.async_increment_pipeline([{"key": "long-window-counter", "increment_value": 1, "ttl": 7200}]) + + ttl_remaining = await cache.async_get_ttl("long-window-counter") + assert ttl_remaining is not None + assert ttl_remaining > time.time() + 600 + + +async def test_async_increment_pipeline_preserves_an_existing_live_ttl_on_later_increments(): + """ + A second increment on the same still-live counter must not reset its + remaining ttl back up to the full window -- only the first increment + (the one that actually creates the counter) should set it. + """ + cache = InMemoryCache(default_ttl=600) + await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}]) + ttl_after_first = cache.ttl_dict["counter"] + + await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}]) + ttl_after_second = cache.ttl_dict["counter"] + + assert ttl_after_second == ttl_after_first + assert cache.get_cache("counter") == 2 + + def test_in_memory_openai_obj_cache(): from openai import OpenAI 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..aea37ecf149 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -0,0 +1,1594 @@ +""" +Unit tests for the global-scope, model-independent tag rate limiter. +""" + +import asyncio +import os +import time +import uuid +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} + + +def _redis_hook(time_controller: TimeController): + from litellm.caching.redis_cache import RedisCache + + redis_host = os.getenv("REDIS_HOST") + redis_port = os.getenv("REDIS_PORT") + if not redis_host or not redis_port: + pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set") + redis_cache = RedisCache(host=redis_host, port=int(redis_port), password=os.getenv("REDIS_PASSWORD")) + dual_cache = DualCache(redis_cache=redis_cache) + return _PROXY_GlobalTagRateLimitsHook(internal_usage_cache=dual_cache, time_provider=time_controller.now), redis_cache + + +# --------------------------------------------------------------------------- +# 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", + ) + + +@pytest.mark.asyncio +async def test_apply_to_models_accounts_when_a_fallback_retry_re_admits_with_a_different_model( + time_controller, monkeypatch +): + """ + Unlike the test above (one admission, a later Router-level fallback + reported only at success time), this simulates + _pre_call_with_fallbacks itself re-running async_pre_call_hook for the + SAME call_id with a DIFFERENT model after some other hook rejected the + original one. The first admission (opus-chain, in apply_to_models scope) + must still get its success-time accounting even though the second + admission (sonnet-chain, out of scope) is the one that actually proceeds + -- overwriting a single "last admitted model" field would silently drop + the spend for the in-scope entry. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "chain_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"}, + call_type="completion", + ) + # _pre_call_with_fallbacks retry: same call_id, fallback model outside apply_to_models. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "sonnet-chain"}, + call_type="completion", + ) + + 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 $12 spend landed in the opus-chain-scoped bucket, so a fresh + # opus-chain request is now over the $10 limit and 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-2"), "model": "opus-chain"}, + call_type="completion", + ) + + +@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 +# 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_concurrency_reservation_released_when_every_fallback_is_exhausted(time_controller, monkeypatch): + """ + When _pre_call_with_fallbacks exhausts every fallback model (another + hook rejects each one) and re-raises, the request never reaches the + real LLM call -- neither async_log_success_event nor + async_log_failure_event, both tied to that call's own wrapper, ever + fires. async_post_call_failure_hook is the only remaining release path + for a reservation this hook already admitted earlier in that same + fallback chain. + """ + 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) + + # This hook admits and reserves the only slot. + 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", + ) + # _pre_call_with_fallbacks eventually gives up (every fallback rejected + # by some other hook) and reports the failure via post_call_failure_hook. + await hook.async_post_call_failure_hook( + request_data={"litellm_call_id": "call-1"}, + original_exception=ProxyRateLimitError( + detail={"error": "some_other_hooks_limit"}, headers={}, rate_limit_type=None + ), + user_api_key_dict=_key(), + ) + + 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 + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_when_disconnect_runs_in_a_forked_task(time_controller, monkeypatch): + """ + common_request_processing.py's real disconnect path runs the release call + inside a task forked via asyncio.create_task AFTER admission already ran + in the parent task (_await_llm_call_cancelling_on_disconnect's `monitor` + task, the streaming generator's own task), not in the same coroutine as + admission the way the test above calls it. A ContextVar-based stash that + only happens to work when both calls share one coroutine has caused a + real leak in this codebase before, so this exercises the actual + parent-then-forked-child topology instead. + """ + 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") + + async def release_in_child_task() -> None: + await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) + + await asyncio.create_task(release_in_child_task()) + + 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_log_success_event_reads_tags_when_top_level_kwargs_carries_a_null_metadata_key( + time_controller, monkeypatch +): + """ + Bugbot finding: some call paths populate a top-level "metadata" (or + "litellm_metadata") key on kwargs (Logging.model_call_details) set to + None, alongside the real, populated dict nested under + kwargs["litellm_params"]. _get_tags_from_request_kwargs checks the + top-level key first; passing it raw kwargs made a present-but-None + top-level key win, reading no tags at all even though + metadata_variable_name correctly resolved to "metadata". + """ + 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) + + kwargs = { + "litellm_call_id": "call-1", + "metadata": None, + "litellm_params": { + "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", + ) + + +# --------------------------------------------------------------------------- +# 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 at async_log_success_event time is Logging.model_call_details: + # metadata/inherited_tags live nested under kwargs["litellm_params"], + # never at the top level -- see + # test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key + # for the same shape. + kwargs = { + "litellm_call_id": "attack-1", + "litellm_params": { + "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", + ) + + +@pytest.mark.asyncio +async def test_rejection_detail_does_not_disclose_the_resolved_tag_value(time_controller, monkeypatch): + """ + tag_value can resolve from inherited_tags (server-assigned key/team/ + project metadata via order_tags_for_identity_resolution), so echoing it + back in the client-facing 429 detail would disclose that identity to + the caller who triggered the rejection. + """ + 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() + + data = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["company_id:secret-internal-name"], "inherited_tags": ["company_id:secret-internal-name"]}, + } + await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=data, call_type="completion") + + with pytest.raises(ProxyRateLimitError) as exc_info: + await hook.async_pre_call_hook( + user_api_key_dict=key, + cache=DualCache(), + data={ + "litellm_call_id": "call-2", + "metadata": { + "tags": ["company_id:secret-internal-name"], + "inherited_tags": ["company_id:secret-internal-name"], + }, + }, + call_type="completion", + ) + + assert "tag_value" not in exc_info.value.detail + + +# --------------------------------------------------------------------------- +# Concurrency ttl refresh -- veria-ai finding: this hook's own atomic checks +# never carried refresh_ttl through to TAG_RL_CHECK_AND_INCR_SCRIPT or +# InMemoryCache.set_cache, even though it shares that script with +# model_based_tag_rate_limits_hook (whose own call sites got fixed first) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + A concurrency bucket isn't epoch-windowed like requests/tokens/dollars -- + its ttl exists only as a crash-safety net for a reservation whose + explicit release never runs -- so a still-active bucket receiving + continuous admissions must keep extending that ttl, or it expires + mid-flight under sustained traffic, silently admitting past the cap. + """ + hook, redis_cache = _redis_hook(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + key = f"{{tag_rl:test:global-ttl-refresh:{uuid.uuid4().hex}}}:inflight" + cache = hook.internal_usage_cache + try: + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = await redis_cache.init_async_client().ttl(key) + assert ttl_after_first_admission > 0 + + await asyncio.sleep(2) + + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = await redis_cache.init_async_client().ttl(key) + assert ttl_after_second_admission >= 2 + finally: + await redis_cache.async_delete_cache(key=key) + + +@pytest.mark.asyncio +async def test_in_memory_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + In-memory mirror of the Redis test above: InMemoryCache.allow_ttl_override + leaves a still-live ttl untouched, so without refresh_ttl reaching + set_cache a concurrency counter's expiry stayed fixed from its first + admission even under sustained traffic. + """ + hook = _make_hook(time_controller) + cache = hook.internal_usage_cache + in_memory_cache = cache.dual_cache.in_memory_cache + key = f"tag_rl:test:global-in-memory-ttl-refresh:{uuid.uuid4().hex}" + + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = in_memory_cache.ttl_dict[key] + + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = in_memory_cache.ttl_dict[key] + + assert ttl_after_second_admission > ttl_after_first_admission + + +@pytest.mark.asyncio +async def test_concurrency_limit_admission_refreshes_ttl_end_to_end(time_controller, monkeypatch): + """ + End-to-end regression through async_pre_call_hook itself (not just the + low-level _check_and_increment_one helper above): a concurrency entry's + bucket key must carry a live ttl after admission, proving refresh_ttl is + actually wired from the classified check through to the atomic batch, + not just present on the helper's own signature. + """ + 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) + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + + in_memory_cache = hook.internal_usage_cache.dual_cache.in_memory_cache + inflight_keys = [key for key in in_memory_cache.ttl_dict if key.endswith(":inflight")] + assert len(inflight_keys) == 1 + assert in_memory_cache.ttl_dict[inflight_keys[0]] > time.time() + + +# --------------------------------------------------------------------------- +# Admission must resolve the authoritative metadata bucket, not the naive +# key-presence check -- veria-ai finding on the sibling model-based hook, +# same vulnerability class in this hook's own admission call site +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_admission_ignores_a_forged_empty_litellm_metadata_key(time_controller, monkeypatch): + """ + get_metadata_variable_name_from_kwargs picks "litellm_metadata" whenever + that key is merely present, regardless of its value. A caller adding an + empty "litellm_metadata" alongside the real, populated "metadata" made + admission read no tags at all, sailing past every configured limit. + """ + 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) + poisoned = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": {}} + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data={**poisoned, "litellm_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={**poisoned, "litellm_call_id": "call-2"}, + call_type="completion", + ) diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py new file mode 100644 index 00000000000..3e34715a606 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -0,0 +1,5472 @@ +""" +Unit tests for tag-scoped token/request/dollar rate limiting. +""" + +import asyncio +import os +import subprocess +import sys +import uuid +from datetime import datetime, timedelta +from types import SimpleNamespace +from typing import Final + +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.model_based_tag_rate_limits_hook import ( + _PENDING_CONCURRENCY_KEYS_FIELD, + _bucket_key, + _build_group_limits, + _build_limits_index, + _ConfiguredLimit, + _current_admission_token, + _extract_team_id, + _inflight_key, + _pending_reservations_cache_key, + _PROXY_ModelBasedTagRateLimitsHook, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + BACKGROUND_TASKS as _BACKGROUND_TASKS, + CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, +) +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + RoutingGroup, + TagRateLimitEntry, + TagRateLimitScope, +) + + +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_limiter(time_controller: TimeController) -> _PROXY_ModelBasedTagRateLimitsHook: + return _PROXY_ModelBasedTagRateLimitsHook( + internal_usage_cache=DualCache(), + time_provider=time_controller.now, + ) + + +def _call_context(tags: list[str]) -> tuple[dict, dict]: + """ + A (request_kwargs, kwargs) pair sharing one `model_call_details` dict, + mirroring production: admission reads `request_kwargs["litellm_logging_obj"] + .model_call_details`, and the `kwargs` passed to async_log_success_event / + async_log_failure_event / async_release_disconnect_state_hook's + request_data *is* that same model_call_details dict (or carries the same + logging_obj) -- see _PENDING_CONCURRENCY_KEYS_FIELD's docstring. A plain + SimpleNamespace stands in for the real Logging object; only its + model_call_details attribute is used. + """ + model_call_details: dict = {} + logging_obj = SimpleNamespace(model_call_details=model_call_details) + request_kwargs = {"metadata": {"tags": tags}, "litellm_logging_obj": logging_obj} + # kwargs must be the *same* dict object model_call_details is, so that + # admission's writes onto model_call_details are visible when this kwargs + # is later passed to a release hook -- see the docstring above. + model_call_details["litellm_logging_obj"] = logging_obj + model_call_details["metadata"] = {"tags": tags} + return request_kwargs, model_call_details + + +def _deployment(model_name: str, deployment_id: str, tag_rate_limits: dict) -> dict: + return { + "model_name": model_name, + "litellm_params": {"model": "gpt-4o", "mock_response": "ok"}, + "model_info": {"id": deployment_id, "tag_rate_limits": tag_rate_limits}, + } + + +def _expected_bucket_key( + model_group: str, + unit: str, + name: str, + tag_id: str, + tag_value: str, + period_seconds: int, + now: float, + deployment_scope: tuple | None = None, + team_scope: str | None = None, + resolved_group: str | None = None, + key_hash: str | None = None, + limit: float = 1, + enabled_for: dict | None = None, + disabled_for: dict | None = None, + scope_by_key_hash: bool = False, +) -> str: + """ + Builds the exact key the real code would compute (via _hash_tag's + fixed-length hashing of tag_value), instead of hand-writing the raw + tag value into a literal string -- the internal key format (hashed or + not) is an implementation detail these tests shouldn't hardcode. + + `limit` and the scoping fields default to values that produce a + stable fingerprint for tests that don't care about it, but must be + passed matching the real entry's own configuration whenever a test's + router declares a `limit` other than 1 (or any scoping) for the entry + whose key this reproduces -- see _policy_fingerprint, which folds them + into the key precisely so two differently-configured entries sharing a + name never collide onto the same counter. + """ + configured = _ConfiguredLimit( + unit=unit, + entry=TagRateLimitEntry( + name=name, + tag_id=tag_id, + limit=limit, + period_seconds=period_seconds, + enabled_for=enabled_for, + disabled_for=disabled_for, + scope_by_key_hash=scope_by_key_hash, + ), + deployment_scope=deployment_scope, + team_scope=team_scope, + resolved_group=resolved_group, + ) + bucket_id = int(now) // period_seconds + return _bucket_key(model_group, configured, tag_value, bucket_id, key_hash=key_hash) + + +# --------------------------------------------------------------------------- +# pending-reservations cache key (async_post_call_failure_hook mirror) +# --------------------------------------------------------------------------- + + +def test_pending_reservations_cache_key_bounds_call_id_regardless_of_input_size(): + """ + veria-ai finding on PR #36541: litellm_call_id comes straight from the + caller-controlled x-litellm-call-id header with no length bound, and was + embedded directly in the pending-reservations mirror key -- a caller + submitting long ids across many in-flight tagged requests could inflate + Redis/in-memory key size disproportionately. Hashed via + _fixed_length_identity, same as every other caller-controlled value this + hook puts in a cache key. + """ + huge_call_id = "x" * 5_000_000 + key = _pending_reservations_cache_key(huge_call_id, "some-key-hash") + assert len(key) < 200 + + +def test_pending_reservations_cache_key_preserves_distinctness(): + assert _pending_reservations_cache_key("call-a", "kh") != _pending_reservations_cache_key("call-b", "kh") + assert _pending_reservations_cache_key("call-a", "kh") == _pending_reservations_cache_key("call-a", "kh") + + +@pytest.mark.asyncio +async def test_an_oversized_tag_value_does_not_inflate_the_bucket_key(time_controller): + """ + End-to-end: a request tagged with a multi-megabyte end_user_id value + must still resolve to a short, fixed-length bucket key, not one whose + size scales with the caller's input. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1000, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + huge_tag_value = "y" * 2_000_000 + + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:{huge_tag_value}"]}}, + ) + + (only_key,) = limiter.internal_usage_cache.dual_cache.in_memory_cache.cache_dict.keys() + assert len(only_key) < 200 + + +# --------------------------------------------------------------------------- +# _extract_team_id -- must read only the one field the server actually +# authenticates into, never fall back to the other +# --------------------------------------------------------------------------- + + +def test_extract_team_id_ignores_a_forged_value_in_the_non_authoritative_field(): + request_kwargs = { + "metadata": {"user_api_key_team_id": "forged-team"}, + "litellm_metadata": {"user_api_key_team_id": "real-team"}, + } + assert _extract_team_id(request_kwargs, "litellm_metadata") == "real-team" + + +@pytest.mark.asyncio +async def test_filter_deployments_ignores_a_forged_empty_litellm_metadata_key(time_controller): + """ + Veria AI finding: get_metadata_variable_name_from_kwargs picks + "litellm_metadata" whenever that key is merely present, regardless of its + value. add_litellm_data_to_request writes real, authenticated team/tag + identity into "metadata" for an ordinary (non LITELLM_METADATA_ROUTES) + request, but leaves any caller-supplied "litellm_metadata" sitting + alongside it untouched -- so a caller adding an empty "litellm_metadata" + to a chat completions request made admission read no tags/team at all, + sailing past every configured limit. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": {}} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_present_but_none(time_controller): + """ + Same misresolution, naturally occurring rather than attacker-forged: a + plain chat completion's kwargs carries a "litellm_metadata" key that is + always present but set to None, alongside the real, populated "metadata" + dict (see resolve_authoritative_metadata_variable_name's own docstring). + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": None} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_ignores_a_forged_populated_litellm_metadata_key(time_controller): + """ + Bugbot finding on the fix above: requiring litellm_metadata to be merely + non-empty is still forgeable -- a caller can populate it with its own, + unrelated keys on an ordinary route, which is non-empty but carries none + of the real identity the proxy wrote into "metadata". Only + add_litellm_data_to_request's own "user_api_key_auth" marker, stripped + from any bucket the caller doesn't own, proves a bucket is authoritative. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "litellm_metadata": {"x": 1}, + } + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + +# --------------------------------------------------------------------------- +# TagRateLimitEntry -- limit validation +# --------------------------------------------------------------------------- + + +def test_tag_rate_limit_entry_rejects_nan_limit(): + """ + NaN compares False against every ordering operator, so a NaN limit makes + the atomic requests/concurrency check-and-increment (rejects when the new + value exceeds the limit) admit indefinitely, while the read-only + tokens/dollars check (admits when the current value is under the limit) + rejects every tagged request -- either way silently defeating the entry. + """ + with pytest.raises(ValidationError, match="limit must not be NaN"): + TagRateLimitEntry(name="n", limit=float("nan"), period_seconds=60) + + +def test_tag_rate_limit_entry_rejects_infinite_limit(): + """ + Positive infinity makes the atomic requests/concurrency + current + increment > limit check always false, so admission never + rejects; negative infinity makes it always true, rejecting every tagged + request. Same silent-misconfiguration class as NaN, just via a different + non-finite float rather than a non-ordering one. + """ + with pytest.raises(ValidationError, match="limit must be finite"): + TagRateLimitEntry(name="n", limit=float("inf"), period_seconds=60) + with pytest.raises(ValidationError, match="limit must be finite"): + TagRateLimitEntry(name="n", limit=float("-inf"), period_seconds=60) + + +def test_tag_rate_limit_entry_rejects_zero_or_negative_limit(): + """ + A limit of 0 (or negative) makes the atomic requests/concurrency check + (current + increment > limit) reject every admission and the read-only + tokens/dollars check (current < limit) never admit, same silent + always-reject-everything failure mode as a negative-infinity limit -- + almost certainly a config typo, not an intentional "block everything" + policy, so reject it at config load time instead. + """ + with pytest.raises(ValidationError, match="limit must be a positive number"): + TagRateLimitEntry(name="n", limit=0, period_seconds=60) + with pytest.raises(ValidationError, match="limit must be a positive number"): + TagRateLimitEntry(name="n", limit=-1, period_seconds=60) + + +# --------------------------------------------------------------------------- +# TagRateLimitEntry -- period_seconds validation +# --------------------------------------------------------------------------- + + +def test_tag_rate_limit_entry_rejects_zero_period_seconds(): + with pytest.raises(ValidationError, match="period_seconds must be a positive integer"): + TagRateLimitEntry(name="n", limit=1, period_seconds=0) + + +def test_tag_rate_limit_entry_rejects_negative_period_seconds(): + with pytest.raises(ValidationError, match="period_seconds must be a positive integer"): + TagRateLimitEntry(name="n", limit=1, period_seconds=-1) + + +def test_tag_rate_limit_entry_accepts_positive_period_seconds(): + entry = TagRateLimitEntry(name="n", limit=1, period_seconds=60) + assert entry.period_seconds == 60 + + +# --------------------------------------------------------------------------- +# _build_group_limits -- chain-wide vs per-deployment scoping +# --------------------------------------------------------------------------- + + +def test_build_group_limits_chain_wide_when_all_deployments_agree(): + deployments = [ + _deployment( + "grp", "dep-1", {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}} + ), + _deployment( + "grp", "dep-2", {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}} + ), + ] + configured = _build_group_limits(deployments, "tokens") + assert len(configured) == 1 + assert configured[0].deployment_scope is None + + +def test_build_group_limits_per_deployment_when_values_diverge(): + """ + Regression test: a naive index that dedupes by (model_name, limit name) and + keeps whichever deployment it encounters first silently drops the second + deployment's config. Divergent values must produce two independent + per-deployment-scoped entries instead. + """ + deployments = [ + _deployment( + "grp", "dep-1", {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}} + ), + _deployment( + "grp", "dep-2", {"token_limits": {"limits": [{"name": "daily", "limit": 999, "period_seconds": 86400}]}} + ), + ] + configured = _build_group_limits(deployments, "tokens") + assert len(configured) == 2 + scopes = {c.deployment_scope for c in configured} + assert scopes == {("dep-1",), ("dep-2",)} + limits = {c.deployment_scope: c.entry.limit for c in configured} + assert limits[("dep-1",)] == 500 + assert limits[("dep-2",)] == 999 + + +def test_build_group_limits_per_deployment_when_only_some_declare_it(): + deployments = [ + _deployment( + "grp", "dep-1", {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}} + ), + _deployment("grp", "dep-2", {}), + ] + configured = _build_group_limits(deployments, "tokens") + assert len(configured) == 1 + assert configured[0].deployment_scope == ("dep-1",) + + +def test_build_group_limits_empty_when_no_deployment_configures_unit(): + deployments = [_deployment("grp", "dep-1", {}), _deployment("grp", "dep-2", {})] + assert _build_group_limits(deployments, "tokens") == () + + +# --------------------------------------------------------------------------- +# TagRateLimitEntry / TagRateLimitScope -- scoping field validation +# --------------------------------------------------------------------------- + + +def test_tag_rate_limit_scope_rejects_empty_values(): + with pytest.raises(ValidationError, match="values must be a non-empty list"): + TagRateLimitScope(tag_id="company_id", values=()) + + +def test_tag_rate_limit_entry_rejects_enabled_for_missing_values(): + with pytest.raises(ValidationError): + TagRateLimitEntry(name="daily", limit=1, period_seconds=60, enabled_for={"tag_id": "company_id"}) + + +def test_tag_rate_limit_scope_normalizes_values_order_and_duplicates(): + scope = TagRateLimitScope(tag_id="company_id", values=("1032", "1001", "1001")) + assert scope.values == ("1001", "1032") + + +def test_tag_rate_limit_entry_rejects_empty_apply_to_key_alias(): + with pytest.raises(ValidationError, match="apply_to_key_alias must be a non-empty list"): + TagRateLimitEntry(name="daily", limit=1, period_seconds=60, apply_to_key_alias=()) + + +def test_tag_rate_limit_entry_normalizes_apply_to_key_alias_order_and_duplicates(): + entry = TagRateLimitEntry( + name="daily", limit=1, period_seconds=60, apply_to_key_alias=("team-b-key", "team-a-key", "team-a-key") + ) + assert entry.apply_to_key_alias == ("team-a-key", "team-b-key") + + +def test_tag_rate_limit_entry_rejects_empty_apply_to_models(): + with pytest.raises(ValidationError, match="apply_to_models must be a non-empty list"): + TagRateLimitEntry(name="daily", limit=1, period_seconds=60, apply_to_models=()) + + +def test_tag_rate_limit_entry_normalizes_apply_to_models_order_and_duplicates(): + entry = TagRateLimitEntry( + name="daily", limit=1, period_seconds=60, apply_to_models=("sonnet-chain", "opus-chain", "opus-chain") + ) + assert entry.apply_to_models == ("opus-chain", "sonnet-chain") + + +# --------------------------------------------------------------------------- +# _hash_tag / _bucket_key -- policy identity folds into the Redis key itself +# --------------------------------------------------------------------------- + + +def test_bucket_key_differs_for_same_named_entries_with_different_limits(): + """ + A plain, unscoped entry and a stricter, scoped override can legitimately + share a `name` (the worked example in the docs uses distinct names, but + nothing in validation requires that) -- resolve_any/_build_group_limits + already treat differing limit/scoping as genuinely distinct policies for + dedup purposes, so the actual counter key must too, or two + differently-configured entries that happen to share a name check and + charge the identical Redis/in-memory bucket. + """ + now = 0.0 + default_key = _expected_bucket_key("grp", "requests", "daily", "end_user_id", "u1", 86400, now, limit=2500) + override_key = _expected_bucket_key( + "grp", + "requests", + "daily", + "end_user_id", + "u1", + 86400, + now, + limit=1, + enabled_for={"tag_id": "company_id", "values": ["1032"]}, + ) + assert default_key != override_key + + +def test_bucket_key_differs_for_same_named_entries_with_different_scoping_only(): + now = 0.0 + excluding_u1 = _expected_bucket_key( + "grp", + "requests", + "daily", + "end_user_id", + "u2", + 86400, + now, + limit=100, + disabled_for={"tag_id": "end_user_id", "values": ["u1"]}, + ) + excluding_u2 = _expected_bucket_key( + "grp", + "requests", + "daily", + "end_user_id", + "u2", + 86400, + now, + limit=100, + disabled_for={"tag_id": "end_user_id", "values": ["u2"]}, + ) + assert excluding_u1 != excluding_u2 + + +def test_bucket_key_differs_for_same_named_entries_diverging_only_on_scope_by_key_hash(): + """ + _DedupSignature already folds scope_by_key_hash into dedup (two + deployments declaring the same name/tag_id but different + scope_by_key_hash become two distinct _ConfiguredLimit entries, not one + merged one), but _policy_fingerprint didn't fold it into the bucket-key + hash. When a request's key_hash resolves to None -- e.g. no virtual key + on the call -- both entries' key_hash-derived suffix is empty too, so an + unscoped entry and a key-hash-scoped entry that otherwise share every + other field collided onto the identical counter, letting one entry's + admission or accounting silently corrupt the other's. + """ + now = 0.0 + unscoped = _expected_bucket_key( + "grp", "requests", "daily", "end_user_id", "u1", 86400, now, limit=100, scope_by_key_hash=False + ) + key_hash_scoped_but_no_key_present = _expected_bucket_key( + "grp", "requests", "daily", "end_user_id", "u1", 86400, now, limit=100, scope_by_key_hash=True, key_hash=None + ) + assert unscoped != key_hash_scoped_but_no_key_present + + +# --------------------------------------------------------------------------- +# _build_group_limits -- scoping fields fold into the dedup signature +# --------------------------------------------------------------------------- + + +def test_build_group_limits_per_deployment_when_disabled_for_diverges(): + """ + Regression test: two deployments agreeing on tag_id/limit/period_seconds + but declaring different disabled_for scopes are genuinely different + policies and must not be silently merged into one shared bucket -- the + same class of bug test_build_group_limits_per_deployment_when_values_diverge + already guards against for a plain divergent limit value. + """ + deployments = [ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "limit": 500, + "period_seconds": 86400, + "disabled_for": {"tag_id": "end_user_id", "values": ["u1"]}, + } + ] + } + }, + ), + _deployment( + "grp", + "dep-2", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "limit": 500, + "period_seconds": 86400, + "disabled_for": {"tag_id": "end_user_id", "values": ["u2"]}, + } + ] + } + }, + ), + ] + configured = _build_group_limits(deployments, "tokens") + assert len(configured) == 2 + scopes = {c.deployment_scope for c in configured} + assert scopes == {("dep-1",), ("dep-2",)} + + +def test_build_group_limits_chain_wide_when_disabled_for_agrees(): + deployments = [ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "limit": 500, + "period_seconds": 86400, + "disabled_for": {"tag_id": "end_user_id", "values": ["u1"]}, + } + ] + } + }, + ), + _deployment( + "grp", + "dep-2", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "limit": 500, + "period_seconds": 86400, + "disabled_for": {"tag_id": "end_user_id", "values": ["u1"]}, + } + ] + } + }, + ), + ] + configured = _build_group_limits(deployments, "tokens") + assert len(configured) == 1 + assert configured[0].deployment_scope is None + + +def test_build_group_limits_chain_wide_when_disabled_for_agrees_in_different_order(): + """ + Two deployments declaring the identical disabled_for values set, just in + a different config order, must dedup to one chain-wide entry -- config + order is not a policy difference. Relies on TagRateLimitScope's own + normalization (sorting) of values at construction time, not on this + dedup path re-sorting them itself. + """ + deployments = [ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "limit": 500, + "period_seconds": 86400, + "disabled_for": {"tag_id": "end_user_id", "values": ["u1", "u2"]}, + } + ] + } + }, + ), + _deployment( + "grp", + "dep-2", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "limit": 500, + "period_seconds": 86400, + "disabled_for": {"tag_id": "end_user_id", "values": ["u2", "u1"]}, + } + ] + } + }, + ), + ] + configured = _build_group_limits(deployments, "tokens") + assert len(configured) == 1 + assert configured[0].deployment_scope is None + + +# --------------------------------------------------------------------------- +# async_filter_deployments -- enforcement +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_filter_deployments_noop_without_config(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router(model_list=[_deployment("grp", "dep-1", {})]) + limiter.update_variables(llm_router=router) + + healthy = router.model_list + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_filter_deployments_allows_under_limit_and_rejects_at_limit(time_controller): + """ + "requests" admission is atomic check-and-increment at the filter step + itself (not a separate read-then-account-later pass), so two concurrent + requests can never both read "1 under limit" and both get admitted past + a limit of 2 -- each call's own increment is immediately visible to the + next. Calling the filter 3 times with limit=2 must admit exactly 2 and + reject the 3rd. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 2, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + for _ in range(2): + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert exc_info.value.status_code == 429 + assert exc_info.value.detail["tag_value"] == "u1" + assert exc_info.value.detail["limit_name"] == "per_minute" + + +@pytest.mark.asyncio +async def test_filter_deployments_falls_back_to_deployment_model_name_for_routing_group_calls(time_controller): + """ + Router keeps a callable routing-group name distinct from every member + deployment's own model_name (see Router._get_routing_group_deployments), + so async_filter_deployments can be called with model="my-group" while + healthy_deployments carries the group's real member deployments. The + limiter must still resolve and enforce each member's own configured + limits rather than silently no-opping because "my-group" itself never + appears in the index. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "backend-a", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 2, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + for _ in range(2): + result = await limiter.async_filter_deployments( + model="my-group", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="my-group", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_routing_group_does_not_collide_across_different_model_names(time_controller): + """ + A routing group can span deployments from different model_names that + happen to declare an identically-named, identically-configured limit. + Each must get its own bucket (keyed by its own model_name via + resolved_group), not share one just because the caller addressed both + through the same group name. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "backend-a", + "dep-a", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ), + _deployment( + "backend-b", + "dep-b", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ), + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # Exhaust backend-a's limit (limit=1) via the group-addressed call. + await limiter.async_filter_deployments( + model="my-group", + healthy_deployments=[healthy[0]], + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="my-group", + healthy_deployments=[healthy[0]], + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # backend-b's own bucket must be untouched -- same group, same tag, same + # limit name, but a different underlying model_name. + result = await limiter.async_filter_deployments( + model="my-group", + healthy_deployments=[healthy[1]], + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == [healthy[1]] + + +def test_resolve_any_dedups_identical_signature_across_member_model_names(): + """ + A real routing-group hop presents every member simultaneously (Router + resolves the group name to its full member list for one filtering pass, + then picks exactly one afterwards), and resolve_any is called once for + that one hop with every member's model_name as a candidate. Two members + declaring the identical concurrency signature must resolve to one shared + entry for that hop, not two: `async_filter_deployments` checks and + atomically increments every entry `resolve_any` returns as belonging to + this one hop, so two entries here means the hop reserves capacity twice + (once per member) even though only one deployment will actually serve -- + over-charging the caller's own usage and risking a false 429 against a + sibling member that was never over its own limit. + + Admission-level round-trip tests can't distinguish this from "two + separate entries with identical limits, always incremented together": + every hop that presents the same member set moves both buckets in + lockstep regardless of whether they're actually one shared entry or two, + so the dedup can only be verified directly at this level. + """ + concurrency_limits = { + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + } + } + index = _build_limits_index( + [ + _deployment("backend-a", "dep-a", concurrency_limits), + _deployment("backend-b", "dep-b", concurrency_limits), + ] + ) + resolved = index.resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b")) + assert len(resolved) == 1 + assert resolved[0].unit == "concurrency" + assert resolved[0].entry.limit == 1 + + +def test_resolve_any_keeps_divergent_signatures_across_member_model_names_separate(): + """ + Companion to the dedup test above: members that genuinely disagree on + the limit for the same tag_id+name must not be silently collapsed -- + which of two different limits would even apply isn't knowable at this + admission-time hook, before a specific deployment is picked, so both + stay as their own entries (today's pre-existing behavior for a + divergent config, left unchanged by the identical-signature dedup). + """ + index = _build_limits_index( + [ + _deployment( + "backend-a", + "dep-a", + { + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + } + }, + ), + _deployment( + "backend-b", + "dep-b", + { + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 2, "period_seconds": 300}] + } + }, + ), + ] + ) + resolved = index.resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b")) + assert len(resolved) == 2 + assert {c.entry.limit for c in resolved} == {1, 2} + + +def test_resolve_any_keeps_divergent_disabled_for_across_member_model_names_separate(): + """ + resolve_any's own dedup key omitted enabled_for/disabled_for/ + apply_to_key_alias, so two routing-group members agreeing on + tag_id/limit/period_seconds but declaring different disabled_for scopes + collapsed to whichever model_name sorted first -- silently applying the + wrong member's policy (and, for the discarded one, no enforcement or + accounting at all for callers only that policy covers). This is the same + class of bug test_build_group_limits_per_deployment_when_disabled_for_diverges + already guards against for the sibling load-balanced-group dedup path. + """ + concurrency_limits_excluding_u1 = { + "concurrency_limits": { + "limits": [ + { + "name": "inflight", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 300, + "disabled_for": {"tag_id": "end_user_id", "values": ["u1"]}, + } + ] + } + } + concurrency_limits_excluding_u2 = { + "concurrency_limits": { + "limits": [ + { + "name": "inflight", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 300, + "disabled_for": {"tag_id": "end_user_id", "values": ["u2"]}, + } + ] + } + } + index = _build_limits_index( + [ + _deployment("backend-a", "dep-a", concurrency_limits_excluding_u1), + _deployment("backend-b", "dep-b", concurrency_limits_excluding_u2), + ] + ) + resolved = index.resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b")) + assert len(resolved) == 2 + assert {c.entry.disabled_for.values for c in resolved} == {("u1",), ("u2",)} + + +def test_resolve_any_keeps_divergent_apply_to_models_across_member_model_names_separate(): + """Same class of bug as the disabled_for test above, for apply_to_models: + two routing-group members agreeing on tag_id/limit/period_seconds but + scoped to different apply_to_models lists must not collapse to one.""" + concurrency_limits_for_opus = { + "concurrency_limits": { + "limits": [ + { + "name": "inflight", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 300, + "apply_to_models": ["opus-chain"], + } + ] + } + } + concurrency_limits_for_sonnet = { + "concurrency_limits": { + "limits": [ + { + "name": "inflight", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 300, + "apply_to_models": ["sonnet-chain"], + } + ] + } + } + index = _build_limits_index( + [ + _deployment("backend-a", "dep-a", concurrency_limits_for_opus), + _deployment("backend-b", "dep-b", concurrency_limits_for_sonnet), + ] + ) + resolved = index.resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b")) + assert len(resolved) == 2 + assert {c.entry.apply_to_models for c in resolved} == {("opus-chain",), ("sonnet-chain",)} + + +def test_resolve_any_picks_the_same_resolved_group_regardless_of_hash_seed(): + """ + Two members with an identical signature dedup to whichever one + `frozenset(candidate_model_names)` iterates first. Plain `frozenset` + iteration order for strings is seeded from `PYTHONHASHSEED`, which is + randomized per process by default, so two proxy worker processes (or the + same process across a restart) resolving the identical member set could + pick different members as `resolved_group` -- fragmenting what's meant to + be one shared Redis bucket into two. This can't be observed from within + one interpreter (a single process has one fixed seed for its lifetime), + so this spawns two real subprocesses pinned to seeds empirically known to + order these three names differently under a plain, unsorted frozenset -- + see the bug report this regression-tests for the exact reproduction. + """ + script = ( + "from litellm.proxy.hooks.model_based_tag_rate_limits_hook import _build_limits_index\n" + "def _deployment(model_name, deployment_id, tag_rate_limits):\n" + " return {'model_name': model_name, 'litellm_params': {'model': 'gpt-4o'}," + " 'model_info': {'id': deployment_id, 'tag_rate_limits': tag_rate_limits}}\n" + "limits = {'concurrency_limits': {'limits': [{'name': 'inflight', 'tag_id': 'end_user_id'," + " 'limit': 1, 'period_seconds': 300}]}}\n" + "index = _build_limits_index([" + "_deployment('backend-a', 'dep-a', limits)," + "_deployment('backend-b', 'dep-b', limits)," + "_deployment('backend-c', 'dep-c', limits)])\n" + "resolved = index.resolve_any('my-group', team_id=None," + " candidate_model_names=('backend-a', 'backend-b', 'backend-c'))\n" + "print(resolved[0].resolved_group)\n" + ) + # seed=1 and seed=3 are empirically confirmed to order these three + # literal strings differently under plain (unsorted) frozenset iteration. + results = { + seed: subprocess.run( + [sys.executable, "-c", script], + env={**os.environ, "PYTHONHASHSEED": seed}, + capture_output=True, + text=True, + check=True, + ).stdout.strip() + for seed in ("1", "3") + } + assert results["1"] == results["3"] == "backend-a" + + +@pytest.mark.asyncio +async def test_filter_deployments_per_entry_fail_open_when_tag_absent(time_controller): + """ + Two entries on the same chain, different tag_ids. Only the tag that's + actually present in the request gets checked; the other has zero effect. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + {"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}, + {"name": "monthly", "tag_id": "team_id", "limit": 1, "period_seconds": 2592000}, + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # Only end_user_id is present -- team_id-keyed entry must not raise or touch Redis. + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + now = time_controller.now().timestamp() + team_bucket_id = int(now) // 2592000 + team_key = f"{{tag_rl:grp:requests:monthly:team_id:chain:whatever}}:{team_bucket_id}" + assert await limiter.internal_usage_cache.async_get_cache(key=team_key, litellm_parent_otel_span=None) is None + + +def _company_tiered_cap_router(default_limit: int, override_limit: int) -> "litellm.Router": + return litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "default_daily", + "tag_id": "end_user_id", + "limit": default_limit, + "period_seconds": 86400, + }, + { + "name": "company_1032_daily", + "tag_id": "end_user_id", + "limit": override_limit, + "period_seconds": 86400, + "enabled_for": {"tag_id": "company_id", "values": ["1032"]}, + "disabled_for": {"tag_id": "end_user_id", "values": ["u1"]}, + }, + ] + } + }, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_scoped_override_skips_for_an_excluded_identity(time_controller): + """ + Company-tiered-cap example from the plan: a stricter override entry + gated to one company via enabled_for, with a handful of named users + excluded from it via disabled_for on the entry's own tag_id. An excluded + user must fall through to the unscoped default entry entirely -- the + override never enforces or accounts for them. + """ + limiter = _make_limiter(time_controller) + router = _company_tiered_cap_router(default_limit=3, override_limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + for _ in range(3): + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1", "company_id:1032"]}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1", "company_id:1032"]}}, + ) + assert exc_info.value.detail["limit_name"] == "default_daily" + + +@pytest.mark.asyncio +async def test_filter_deployments_scoped_override_enforces_for_a_non_excluded_identity_in_scope(time_controller): + """ + The same override applies, and enforces its own stricter limit, for a + company-1032 user who is not disabled_for's excluded identity, proving + the two entries are independently enforced rather than one silently + replacing the other. + """ + limiter = _make_limiter(time_controller) + router = _company_tiered_cap_router(default_limit=3, override_limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u2", "company_id:1032"]}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u2", "company_id:1032"]}}, + ) + assert exc_info.value.detail["limit_name"] == "company_1032_daily" + + +@pytest.mark.asyncio +async def test_filter_deployments_scoped_override_does_not_apply_outside_its_enabled_for_gate(time_controller): + """A user not tagged with the gate company at all only ever hits the + unscoped default entry, even though the override's own limit is looser + and would otherwise still have room.""" + limiter = _make_limiter(time_controller) + router = _company_tiered_cap_router(default_limit=1, override_limit=5) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u3"]}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u3"]}}, + ) + assert exc_info.value.detail["limit_name"] == "default_daily" + + +@pytest.mark.asyncio +async def test_different_tag_ids_with_same_name_do_not_share_a_counter(time_controller): + """ + Regression test: the key format used to omit `tag_id`, so two + independently configured entries sharing the same unit/name (here both + named "daily") but keyed on different tag_ids would collide whenever a + caller's value for one tag_id happened to equal another caller's value + for the other tag_id. With equal limits of 1, a colliding shared counter + would make team_id "u1"'s very first request get wrongly rejected right + after end_user_id "u1"'s own first (and separately limited) request -- + the failure mode a higher team_id limit would have masked. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + {"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}, + {"name": "daily", "tag_id": "team_id", "limit": 1, "period_seconds": 86400}, + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # end_user_id "u1" makes its one allowed request. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # team_id "u1" -- identical value, different tag_id, its own untouched + # limit of 1 -- must still admit its first request. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["team_id:u1"]}} + ) + assert result == healthy + + # Both identities are now genuinely at their own limit of 1. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["team_id:u1"]}}, + ) + + +# --------------------------------------------------------------------------- +# Load-balanced group: chain-wide vs per-deployment enforcement end to end +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_load_balanced_group_per_deployment_breach_rejects_whole_hop(time_controller): + """ + Two deployments share one model_name with divergent per-deployment + limits. Breaching one deployment's bucket rejects the hop even though the + other deployment (still present in healthy_deployments) is comfortably + under its own limit -- this does NOT filter the breaching deployment out + and let the router retry the sibling. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ), + _deployment( + "grp", + "dep-2", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 999, "period_seconds": 86400}] + } + }, + ), + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}} + + now = time_controller.now().timestamp() + dep1_key = _expected_bucket_key( + "grp", "requests", "daily", "end_user_id", "u1", 86400, now, deployment_scope=("dep-1",) + ) + await limiter.internal_usage_cache.async_set_cache(key=dep1_key, value=1, ttl=86400, litellm_parent_otel_span=None) + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + +# --------------------------------------------------------------------------- +# async_log_success_event -- accounting +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_log_success_event_increments_configured_units(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + }, + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 100, "period_seconds": 86400}] + }, + "dollar_limits": { + "limits": [ + {"name": "monthly", "tag_id": "end_user_id", "limit": 50.0, "period_seconds": 2592000} + ] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + # accounting is fired via asyncio.create_task; let it run. + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + dollar_key = _expected_bucket_key("grp", "dollars", "monthly", "end_user_id", "u1", 2592000, now, limit=50.0) + + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=dollar_key, litellm_parent_otel_span=None)) == 0.01 + ) + + # "requests" is accounted atomically at admission (async_filter_deployments), + # not here -- async_log_success_event must not touch its bucket at all. + request_key = _expected_bucket_key("grp", "requests", "daily", "end_user_id", "u1", 86400, now, limit=100) + assert await limiter.internal_usage_cache.async_get_cache(key=request_key, litellm_parent_otel_span=None) is None + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key(time_controller): + """ + 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 accounting for + this route shape. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "litellm_metadata": None, + "metadata": {"tags": ["end_user_id:u1"]}, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + +@pytest.mark.asyncio +async def test_log_success_event_ignores_a_forged_populated_litellm_metadata_key(time_controller): + """ + Same misresolution as the null-key case above, but with a populated + (not merely present) forged litellm_metadata -- a caller-supplied dict + with unrelated keys is still not the field the proxy wrote identity + into. Only the unconditionally-stamped "user_api_key_auth" marker, + which litellm_pre_call_utils.py strips from any bucket a caller doesn't + own, proves a bucket is authoritative. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "litellm_metadata": {"x": 1}, + "metadata": {"tags": ["end_user_id:u1"]}, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_the_key_backed_tag_not_a_caller_forged_one(time_controller): + """ + Bugbot finding: admission (async_filter_deployments) sees a flat + request_kwargs where metadata.inherited_tags sits at the top level, but + kwargs at async_log_success_event time is Logging.model_call_details, + which only ever nests metadata under kwargs["litellm_params"] (see + test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key). + order_tags_for_identity_resolution's own inherited_tags lookup only + checked the top level, so a caller-forged company_id tag that admission + correctly ignored could still get accounted against at success time, + charging a different bucket than the one admission actually checked. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "company_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "metadata": { + "tags": ["company_id:attacker-chosen"], + "inherited_tags": ["company_id:real-company"], + }, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + real_key = _expected_bucket_key("grp", "tokens", "daily", "company_id", "real-company", 86400, now, limit=500000) + forged_key = _expected_bucket_key( + "grp", "tokens", "daily", "company_id", "attacker-chosen", 86400, now, limit=500000 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=real_key, litellm_parent_otel_span=None)) == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=forged_key, litellm_parent_otel_span=None) is None + + +@pytest.mark.asyncio +async def test_log_success_event_reads_nested_litellm_metadata_when_that_is_authoritative(time_controller): + """ + kwargs here is Logging.model_call_details: on LITELLM_METADATA_ROUTES + (/v1/messages, /responses, ...) metadata/litellm_metadata are never + top-level keys, only nested under kwargs["litellm_params"] -- and the + caller's own native "metadata" can be present there with no tags at all, + while the real, server-computed tags live in "litellm_metadata". + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "metadata": {"tags": []}, + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_auth": {}}, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + +@pytest.mark.asyncio +async def test_log_success_event_falls_back_to_serving_deployment_model_name_for_routing_group_calls( + time_controller, +): + """ + standard_logging_object["model_group"] is the caller-visible name from + Router._update_kwargs_before_fallbacks -- for a routing-group call this + is the group name too, which never appears in the index. Success + accounting must fall back to the model_name of the deployment that + actually served this hop (standard_logging_object["model_id"]). + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "backend-a", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "my-group", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("backend-a", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_against_the_same_bucket_admission_checked(time_controller): + """ + resolve_any dedups an identical signature across a routing group's + members into one shared entry, stamped with resolved_group from + whichever member frozenset(candidate_model_names) yields first (see + resolve_any's own docstring). Success accounting for tokens/dollars only + learns the one deployment that actually served this hop; passing just + that single name as resolve_any's sole candidate would make its dedup + trivially resolve to that deployment's own name -- which can differ from + whichever member admission's full-group view picked, silently + accounting usage against a bucket admission never checked and letting a + token/dollar limit be bypassed. Success accounting must reconstruct the + full routing-group candidate set so it lands on the identical bucket + regardless of which member actually served. + """ + token_limits = { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + } + router = litellm.Router( + model_list=[ + _deployment("backend-a", "dep-a", token_limits), + _deployment("backend-b", "dep-b", token_limits), + ], + routing_groups=[ + RoutingGroup(group_name="my-group", models=["backend-a", "backend-b"], routing_strategy="simple-shuffle") + ], + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + + # What admission would check: it sees every member, and resolve_any's + # dedup picks whichever one frozenset yields first for the shared entry. + admitted = limiter._index.get(router).resolve_any( + "my-group", team_id=None, candidate_model_names=("backend-a", "backend-b") + ) + assert len(admitted) == 1 + admission_bucket_group = admitted[0].resolved_group + + # Force the deployment that actually serves to be the *other* member -- + # deterministic regardless of which one frozenset happened to pick above, + # so this test always exercises the mismatch the fix guards against. + serving_model_name = "backend-b" if admission_bucket_group == "backend-a" else "backend-a" + serving_deployment_id = "dep-b" if serving_model_name == "backend-b" else "dep-a" + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "my-group", + "model_id": serving_deployment_id, + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key( + "my-group", + "tokens", + "daily", + "end_user_id", + "u1", + 86400, + now, + resolved_group=admission_bucket_group, + limit=500000, + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + +@pytest.mark.asyncio +async def test_log_success_event_uses_admissions_own_candidate_set_when_group_membership_drifts(time_controller): + """ + Bugbot finding: resolve_any's dedup picks the alphabetically first member + model_name sharing a signature as resolved_group, a pure function of + candidate_model_names. Both admission and success independently rebuild + that set from the router's *live* routing-group membership, so a + deployment added mid-request (a hot-reload) whose name sorts earlier can + make success pick a different resolved_group than admission did, + accounting real usage into a bucket admission never checked. + """ + token_limits = { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + } + router = litellm.Router( + model_list=[ + _deployment("backend-a", "dep-a", token_limits), + _deployment("backend-b", "dep-b", token_limits), + ], + routing_groups=[ + RoutingGroup(group_name="my-group", models=["backend-a", "backend-b"], routing_strategy="simple-shuffle") + ], + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + request_kwargs, model_call_details = _call_context(["end_user_id:u1"]) + + healthy = router._get_routing_group_deployments(model="my-group", team_id=None) + assert healthy is not None + admitted = await limiter.async_filter_deployments( + model="my-group", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert admitted == healthy + admission_bucket_group = ( + limiter._index.get(router) + .resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b"))[0] + .resolved_group + ) + + routing_group = router.get_routing_group("my-group") + assert routing_group is not None + routing_group.models.append("backend-0") + router.add_deployment( + Deployment( + model_name="backend-0", + litellm_params=LiteLLM_Params(model="gpt-4o", mock_response="ok"), + model_info=ModelInfo(id="dep-0", tag_rate_limits=token_limits), + ) + ) + serving_deployment_id = "dep-b" if admission_bucket_group == "backend-a" else "dep-a" + + model_call_details["standard_logging_object"] = { + "model_group": "my-group", + "model_id": serving_deployment_id, + "total_tokens": 42, + "response_cost": 0.01, + } + await limiter.async_log_success_event(kwargs=model_call_details, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + admission_key = _expected_bucket_key( + "my-group", + "tokens", + "daily", + "end_user_id", + "u1", + 86400, + now, + resolved_group=admission_bucket_group, + limit=500000, + ) + drifted_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-0", limit=500000 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=admission_key, litellm_parent_otel_span=None)) + == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=drifted_key, litellm_parent_otel_span=None) is None + + +@pytest.mark.asyncio +async def test_admission_dedups_against_the_full_group_not_just_currently_healthy_members(time_controller): + """ + `healthy_deployments` is Router's cooldown-filtered list for this one + hop -- a member merely cooled down right now is excluded from it, but + it's still a real member of the routing group. Deriving resolve_any's + candidate set from `healthy_deployments` instead of the full group would + make admission's resolved_group choice depend on which members happen to + be healthy at that exact moment, while success accounting (which has no + way to know what was healthy at admission time) always reconstructs the + full, static membership -- landing the two sides on different buckets + whenever a member is cooled down. Admission must dedup against the same + full membership success does, regardless of which members are currently + healthy. + """ + token_limits = { + "token_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 10, "period_seconds": 86400}]} + } + router = litellm.Router( + model_list=[ + _deployment("backend-a", "dep-a", token_limits), + _deployment("backend-b", "dep-b", token_limits), + ], + routing_groups=[ + RoutingGroup(group_name="my-group", models=["backend-a", "backend-b"], routing_strategy="simple-shuffle") + ], + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + + # The shared entry always dedups to "backend-a" (alphabetically first). + # Pre-load *that* bucket over the limit; the "backend-b" bucket (what a + # healthy_deployments-derived candidate set would wrongly resolve to, + # since backend-a is the only one excluded below) stays empty. + now = time_controller.now().timestamp() + over_limit_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-a", limit=10 + ) + await limiter.internal_usage_cache.async_set_cache(key=over_limit_key, value=20.0, litellm_parent_otel_span=None) + + # Simulate backend-a being cooled down: Router would exclude it from the + # healthy_deployments list passed to this hop's admission. + healthy_excluding_backend_a = [d for d in router.model_list if d["model_name"] == "backend-b"] + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="my-group", + healthy_deployments=healthy_excluding_backend_a, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_against_the_key_hash_admission_checked(time_controller): + """ + Admission's `_extract_key_hash` reads `metadata.user_api_key` unconditionally + whenever scope_by_key_hash is set -- that field is already the hashed + token by the time it reaches this hook (see the function's own + docstring), regardless of its shape. `standard_logging_object.metadata` + only ever carries the derived `user_api_key_hash` field, and only when the + raw value happens to look like a SHA-256 hex digest (see + litellm_logging.py's get_standard_logging_metadata) -- a virtual key + represented any other way makes that field silently absent, so reading it + on the success side would account against key_hash=None while admission + scoped the check against the real value, letting usage silently bypass a + per-key limit whenever the key's own representation isn't SHA-256-shaped. + """ + token_limits = { + "token_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 500000, + "period_seconds": 86400, + "scope_by_key_hash": True, + } + ] + } + } + router = litellm.Router(model_list=[_deployment("grp", "dep-1", token_limits)]) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + + # "keyA" deliberately isn't SHA-256-shaped, so standard_logging_object's + # own redaction/derivation step would never populate user_api_key_hash + # for it -- it's simply absent, matching production for a key hash that + # doesn't pass that shape check. + kwargs = { + "metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + "metadata": {}, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + keyed_bucket = _expected_bucket_key( + "grp", "tokens", "daily", "end_user_id", "u1", 86400, now, key_hash="keyA", limit=500000, scope_by_key_hash=True + ) + unkeyed_bucket = _expected_bucket_key( + "grp", "tokens", "daily", "end_user_id", "u1", 86400, now, key_hash=None, limit=500000, scope_by_key_hash=True + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=keyed_bucket, litellm_parent_otel_span=None)) + == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=unkeyed_bucket, litellm_parent_otel_span=None) is None + + +@pytest.mark.asyncio +async def test_log_success_event_charges_the_window_admission_checked_not_a_later_one(time_controller): + """ + Admission classifies its bucket as int(now) // period_seconds at filter + time; success accounting used to recompute a fresh now of its own, so a + call slow enough to cross a period_seconds boundary between admission and + completion got admitted against one window's (still-open) counter but + charged into the next window's fresh, empty one -- silently bypassing the + limit for calls straddling each rollover. Success must charge the exact + window admission classified against, not whatever window happens to be + current when the response finishes. + """ + token_limits = { + "token_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 500, "period_seconds": 60}] + } + } + router = litellm.Router(model_list=[_deployment("grp", "dep-1", token_limits)]) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + + admission_time = time_controller.now().timestamp() + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + # The response takes long enough to cross into the next 60s window before + # completing. + time_controller.advance(61) + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + admitted_window_bucket = _expected_bucket_key( + "grp", "tokens", "per_minute", "end_user_id", "u1", 60, admission_time, limit=500 + ) + later_window_bucket = _expected_bucket_key( + "grp", "tokens", "per_minute", "end_user_id", "u1", 60, time_controller.now().timestamp(), limit=500 + ) + assert ( + float( + await limiter.internal_usage_cache.async_get_cache( + key=admitted_window_bucket, litellm_parent_otel_span=None + ) + ) + == 42.0 + ) + assert ( + await limiter.internal_usage_cache.async_get_cache(key=later_window_bucket, litellm_parent_otel_span=None) + is None + ) + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_against_the_team_id_admission_checked(time_controller): + """ + Admission resolves team_id via `_extract_team_id`, the single + metadata_variable_name-authoritative field lookup -- success used to read + `standard_logging_object.metadata.user_api_key_team_id` instead, a + separately-constructed field that isn't guaranteed to come from the same + field admission used (e.g. on LITELLM_METADATA_ROUTES, where + `litellm_metadata` is authoritative but `standard_logging_object` may + still reflect a different resolution). A mismatched team_id changes + team_scope, which is hashed into the bucket key, so success would charge + a different bucket than the one admission's team-aliased lookup checked. + """ + deployment = _deployment( + "real-model-name", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500, "period_seconds": 86400}] + } + }, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + router = litellm.Router(model_list=[deployment]) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # LITELLM_METADATA_ROUTES shape: litellm_metadata is the authoritative + # field, and team-alias resolution requires the real team_id from it. + request_kwargs = { + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1", "user_api_key_auth": {}} + } + result = await limiter.async_filter_deployments( + model="team-alias-name", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + # standard_logging_object's own team_id deliberately disagrees with the + # real one in litellm_params.litellm_metadata, simulating litellm_logging.py + # resolving a different field than the one admission used. + kwargs = { + "litellm_params": { + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1", "user_api_key_auth": {}} + }, + "standard_logging_object": { + "model_group": "team-alias-name", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + "metadata": {"user_api_key_team_id": None}, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + correct_bucket = _expected_bucket_key( + "team-alias-name", "tokens", "daily", "end_user_id", "u1", 86400, now, team_scope="team-1", limit=500 + ) + wrong_bucket = _expected_bucket_key( + "team-alias-name", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=correct_bucket, litellm_parent_otel_span=None)) + == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=wrong_bucket, litellm_parent_otel_span=None) is None + + +# --------------------------------------------------------------------------- +# concurrency limits -- reserve at admission, release on success/failure +# --------------------------------------------------------------------------- + + +def _concurrency_router(limit: int) -> "litellm.Router": + return litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300}] + } + }, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_cross_unit_rejection_does_not_leave_a_phantom_increment(time_controller): + """ + Regression test: a chain with BOTH a requests limit and a concurrency + limit on the same tag checks both atomically in one + async_filter_deployments call. If the concurrency check rejects the hop, + the requests-unit check (evaluated in the same call, and which would have + been admitted on its own) must NOT have incremented its counter -- + otherwise a rejected hop silently burns through the caller's requests + budget for a call that never actually went through. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 10, "period_seconds": 60}] + }, + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # Occupy the one concurrency slot. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # A second attempt: requests-unit alone would admit (well under 10), but + # concurrency is exhausted, so the whole hop must reject -- and the + # requests counter must remain untouched by this rejected attempt. + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert exc_info.value.detail["type"] == "concurrency" + + now = time_controller.now().timestamp() + request_key = _expected_bucket_key("grp", "requests", "per_minute", "end_user_id", "u1", 60, now, limit=10) + requests_value = await limiter.internal_usage_cache.async_get_cache(key=request_key, litellm_parent_otel_span=None) + assert (float(requests_value) if requests_value is not None else 0.0) == 1.0 + + +@pytest.mark.asyncio +async def test_concurrency_limit_rejects_third_concurrent_reservation(time_controller): + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=2) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + kwargs_1 = {"metadata": {"tags": ["end_user_id:u1"]}} + kwargs_2 = {"metadata": {"tags": ["end_user_id:u1"]}} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=kwargs_1 + ) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=kwargs_2 + ) + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert exc_info.value.detail["type"] == "concurrency" + + +@pytest.mark.asyncio +async def test_requests_admission_is_race_free_under_genuine_concurrency(time_controller): + """ + The concrete race a read-then-account-later design would allow: N + coroutines all read "under limit" before any of them increments, and all + N get admitted even past the limit. Firing many concurrent filter calls + at once (asyncio.gather, not sequential awaits) must admit exactly + `limit`, never more -- the atomic check-and-increment closes the window + a plain GET-then-SET could not. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 5, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + async def attempt(): + try: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + return True + except ProxyRateLimitError: + return False + + results = await asyncio.gather(*(attempt() for _ in range(20))) + assert sum(results) == 5 + + +@pytest.mark.asyncio +async def test_index_refreshes_after_ttl_for_length_preserving_update(time_controller): + """ + Editing an existing deployment's tag_rate_limits in place (same + len(model_list), so the (id(router), len) staleness check alone can't + detect it) must eventually be picked up -- bounded by the index TTL, + not indefinitely stale. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # Same length, deployment mutated in place -- raise the limit to 100. + router.model_list[0]["model_info"]["tag_rate_limits"] = { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 100, "period_seconds": 86400}] + } + } + + time_controller.advance(6) # past _INDEX_TTL_SECONDS + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=router.model_list, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == router.model_list + + +@pytest.mark.asyncio +async def test_concurrency_slot_released_on_success_frees_capacity(time_controller): + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # At capacity: a second concurrent request is rejected. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # The first request completes -- its slot is released -- freeing capacity again. + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 0, + "response_cost": 0, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_background_release_tasks_registry_holds_a_reference_until_done(): + """ + async_log_success_event fires its release via a bare asyncio.create_task + (unlike failure/disconnect, which await it directly) to keep the hot + success path from waiting on a Redis round trip. asyncio.create_task's + own docs warn the event loop only holds a *weak* reference to a task, so + one with no other referrer can be garbage collected before it runs -- + and by the time it would run here, its keys are already popped out of + model_call_details, so a collected task's release is unrecoverable, not + merely delayed. _BACKGROUND_TASKS exists to hold a strong + reference for exactly as long as the task is pending, then release it via + the task's own done-callback -- exercised directly here (an Event gate + gives a deterministic pending window; going through the real + async_log_success_event doesn't, since its own further awaits let a fast + in-memory release resolve before a test could ever observe it pending). + """ + assert len(_BACKGROUND_TASKS) == 0 + gate = asyncio.Event() + + async def _pending_release(): + await gate.wait() + + task = asyncio.create_task(_pending_release()) + _BACKGROUND_TASKS.add(task) + task.add_done_callback(_BACKGROUND_TASKS.discard) + + assert task in _BACKGROUND_TASKS + + gate.set() + await task + + # The done-callback removes it -- the registry doesn't grow unbounded + # across requests. + assert task not in _BACKGROUND_TASKS + assert len(_BACKGROUND_TASKS) == 0 + + +@pytest.mark.asyncio +async def test_success_event_release_is_wired_through_the_background_registry(time_controller): + """ + End-to-end check that async_log_success_event's fire-and-forget release + is genuinely wired through _BACKGROUND_TASKS, not a bare + unreferenced asyncio.create_task -- the registry must be empty again once + the (fast, in-memory) release has had a chance to run, and the release + itself must have actually happened. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 0, + "response_cost": 0, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # The registry was actually populated: proves the release ran through + # _BACKGROUND_TASKS, not a bare unreferenced asyncio.create_task + # (which would never touch this set at all, and an "empty at the end" + # check alone can't tell the two apart -- an empty registry throughout + # would satisfy that just as well as one that filled and drained). + assert len(_BACKGROUND_TASKS) == 1 + + # Two ticks: one for the release task itself to finish (it may already be + # done by the time async_log_success_event returns, given that method's + # own further awaits), and one for its done-callback -- scheduled via + # call_soon when the task completes -- to actually run and discard it. + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert len(_BACKGROUND_TASKS) == 0 + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_success_event_token_accounting_is_wired_through_the_background_registry(time_controller): + """ + Same gap as the concurrency release above, in a second fire-and-forget + task on the same success path: token/dollar accounting is also fired + via a bare asyncio.create_task per cache partition, with no strong + reference of its own. A collected task here drops a usage increment + that can never be recovered (the figures it needed only exist in that + task's own closure), silently under-counting a caller's token/dollar + usage against its configured limit. Must be tracked the same way. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # The registry was actually populated: proves accounting ran through + # _BACKGROUND_TASKS, not a bare unreferenced asyncio.create_task. + assert len(_BACKGROUND_TASKS) == 1 + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert len(_BACKGROUND_TASKS) == 0 + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + +@pytest.mark.asyncio +async def test_concurrency_slot_released_on_disconnect_frees_capacity(time_controller): + """ + A client disconnecting before the first streamed chunk raises + CancelledError/GeneratorExit, which bypasses both async_log_success_event + and async_log_failure_event entirely -- neither fires, so the reservation + would otherwise sit held until the safety-net TTL. The proxy's disconnect + cleanup calls async_release_disconnect_state_hook instead in that case. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # At capacity: a second concurrent request is rejected. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # The first request's client disconnects -- neither logging callback fires -- + # but the disconnect hook still releases its slot, freeing capacity again. + await limiter.async_release_disconnect_state_hook(request_kwargs) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_concurrency_slot_released_on_failure_frees_capacity(time_controller): + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + + kwargs["standard_logging_object"] = {"model_group": "grp"} + await limiter.async_log_failure_event( + kwargs=kwargs, + response_obj=None, + start_time=0, + end_time=0, + ) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_concurrency_slot_released_by_post_call_failure_hook_on_the_final_fallback_hop(time_controller): + """ + litellm's Logging object sets has_logged_async_failure=True after the + first hop's failure and blocks async_log_failure_event for every later + hop (see fallback_event_handlers.py), so a fallback chain's own final, + chain-exhausting failure never reaches async_log_failure_event at all -- + _release_stale_hop_reservations only cleans up a stale reservation when + a *next* hop's admission runs, and there is no next hop after the last + one. async_post_call_failure_hook fires exactly once, at the point the + proxy gives up and returns an error to the caller, regardless of how + many hops ran or whether the completion-level callback was suppressed -- + it must release whatever reservation is still pending at that point. + + request_data here is a distinct dict object from admission's own + request_kwargs, with no litellm_logging_obj at all: proxy/utils.py's + post_call_failure_hook pops that key off request_data before invoking + any callback ("Remove before callbacks iterate — not serialisable"), + and confirmed live, request_data is a third, unrelated object from + every hop's own request_kwargs by the time this fires. litellm_call_id + is the only identifier stable across all of them. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # This hop's admission reserves the slot; its own failure is the chain's + # final one, so async_log_failure_event never fires for it (simulating + # litellm's has_logged_async_failure dedup blocking the callback here). + request_kwargs = { + "metadata": {"tags": ["end_user_id:u1"], "user_api_key": "hash"}, + "litellm_call_id": "call-final", + } + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + await limiter.async_post_call_failure_hook( + request_data={"litellm_call_id": "call-final"}, + original_exception=Exception("all deployments failed"), + user_api_key_dict=UserAPIKeyAuth(api_key="hash"), + ) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_cannot_release_a_different_keys_reservation(time_controller): + """ + Security regression: litellm_call_id comes from the caller-controlled + x-litellm-call-id header, so two different callers choosing the identical + id must not be able to release each other's reservation through the + pending-reservations cache mirror. Request A (key-a, tag victim_user) and + request B (key-b, tag attacker_user) share one call_id; A's own terminal + failure must only ever be able to find and release A's own mirror entry, + keyed by A's server-authenticated key hash, never B's. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + victim_request_kwargs = { + "metadata": {"tags": ["end_user_id:victim_user"], "user_api_key": "key-a-hash"}, + "litellm_call_id": "shared-call-id", + } + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=victim_request_kwargs + ) + + attacker_request_kwargs = { + "metadata": {"tags": ["end_user_id:attacker_user"], "user_api_key": "key-b-hash"}, + "litellm_call_id": "shared-call-id", + } + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=attacker_request_kwargs + ) + + # Simulates request A's own fallback chain exhausting -- its terminal + # failure hook must not touch request B's still-live reservation just + # because both requests share a caller-chosen call_id. + await limiter.async_post_call_failure_hook( + request_data={"litellm_call_id": "shared-call-id"}, + original_exception=Exception("all deployments failed"), + user_api_key_dict=UserAPIKeyAuth(api_key="key-a-hash"), + ) + + # attacker_user's own reservation must still be held: key-a's failure + # hook releasing it would let key-b bypass its own concurrency cap. + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:attacker_user"], "user_api_key": "key-b-hash"}}, + ) + assert exc_info.value.detail["type"] == "concurrency" + + # victim_user's own slot was correctly released by its own key's + # failure hook -- the legitimate single-key path still works. + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:victim_user"], "user_api_key": "key-a-hash"}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_mirror_write_does_not_erase_a_concurrent_siblings_own_reservation(time_controller): + """ + Cursor Bugbot finding: Router.abatch_completion's branches share one + litellm_call_id and one authenticated key, so they mirror their own + concurrency reservation under the identical (call_id, key_hash) cache + entry. Mirroring used to overwrite rather than merge, so branch two's own + admission silently erased branch one's already-mirrored reservation -- + branch one's own eventual chain-exhausting failure would then find + nothing to release there, permanently leaking its slot (until the safety + TTL) even though both branches genuinely finished. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=2) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + shared_call_id = "batch-call-id" + shared_tags = {"tags": ["end_user_id:shared-user"], "user_api_key": "batch-key-hash"} + branch_one_kwargs = {"metadata": shared_tags, "litellm_call_id": shared_call_id} + branch_two_kwargs = {"metadata": shared_tags, "litellm_call_id": shared_call_id} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=branch_one_kwargs + ) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=branch_two_kwargs + ) + + # Both branches finish (order doesn't matter -- whichever hook reads the + # mirror first releases whatever is genuinely still there for either). + await limiter.async_post_call_failure_hook( + request_data={"litellm_call_id": shared_call_id}, + original_exception=Exception("branch one exhausted its fallbacks"), + user_api_key_dict=UserAPIKeyAuth(api_key="batch-key-hash"), + ) + await limiter.async_post_call_failure_hook( + request_data={"litellm_call_id": shared_call_id}, + original_exception=Exception("branch two exhausted its fallbacks"), + user_api_key_dict=UserAPIKeyAuth(api_key="batch-key-hash"), + ) + + # Neither slot leaked: two fresh admissions fit under the limit of 2, + # and a third does not. Under the erasure bug, branch two's write wiped + # branch one's mirrored entry before it was ever read, so only one of + # the two reservations was ever actually released. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared-user"]}}, + ) + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared-user"]}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared-user"]}}, + ) + + +@pytest.mark.asyncio +async def test_concurrency_slot_released_when_a_different_hook_rejects_the_request(time_controller): + """ + global_tag_rate_limits_hook raises the identical ProxyRateLimitError + shape (detail["error"] == "tag_rate_limit_exceeded") this hook's own + admission does. 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 based on the + shared marker alone would leak this hook's own slot until the safety + TTL, even though nothing about this hook's own admission failed. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + + 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="grp", + llm_provider="litellm_proxy", + ) + kwargs["exception"] = other_hooks_rejection + await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_concurrency_slot_released_on_fallback_recovered_hop_failure(time_controller): + """ + A hop that fails but gets recovered by a later fallback never reaches + a terminal, request-level hook -- there is none, deliberately, since + there's no reliable, caller-uncontrolled way to correlate multiple hops + of one logical request from inside a CustomLogger hook (see + async_log_failure_event's docstring for why litellm_call_id, the one + candidate, can't be trusted for this). async_log_failure_event fires + per hop, on every failure, recomputing this hop's own key independently, + so this specific case -- exactly one prior failure, then a fallback that + succeeds -- is still handled correctly. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + kwargs["standard_logging_object"] = {"model_group": "grp", "model_id": "dep-1"} + await limiter.async_log_failure_event( + kwargs=kwargs, + response_obj=None, + start_time=0, + end_time=0, + ) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_pending_concurrency_reservations_do_not_leak_across_unrelated_requests(time_controller): + """ + Security regression test, current design: pending concurrency keys are + stashed on the admitting request's own `model_call_details` dict (see + `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring), never in a registry keyed + by anything caller-visible or by ambient asyncio context. Two unrelated + concurrent requests each get their own `model_call_details` in + production, so one request's release can never see or drain a different + request's still-pending reservation, regardless of which asyncio task + each happens to run in and even when both share the identical tag value + (an earlier design keyed reservations by `litellm_call_id` -- settable by + the caller via the `x-litellm-call-id` header -- which let two unrelated + requests merge reservations simply by choosing the same id). + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_a, kwargs_a = _call_context(["end_user_id:shared"]) + request_b, kwargs_b = _call_context(["end_user_id:shared"]) + + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_a + ) + # B shares A's tag value but is a genuinely separate request/object: at + # capacity (limit=1), B is rejected and never reserves anything. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_b + ) + + # B's own failure event releases via its own (empty) model_call_details -- + # this must not accidentally drain A's still-live reservation. + kwargs_b["standard_logging_object"] = {"model_group": "grp", "model_id": "dep-1"} + await limiter.async_log_failure_event(kwargs=kwargs_b, response_obj=None, start_time=0, end_time=0) + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared"]}}, + ) + + # A's own success event correctly releases its own reservation. + kwargs_a["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 0, + "response_cost": 0, + } + await limiter.async_log_success_event(kwargs=kwargs_a, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_concurrency_released_for_every_hop_across_a_real_task_boundary(time_controller): + """ + litellm dedupes `async_log_failure_event` to fire once per logical + request: only the first failed hop's failure reaches it (see + `Logging.has_run_logging`'s `has_logged_async_failure` guard); a later + failed hop (a retry or a further fallback) never gets its own failure + event at all. Reservations still accumulate at admission for every hop + regardless (onto the request's own `model_call_details`, shared across + every hop of one logical request -- see `_PENDING_CONCURRENCY_KEYS_FIELD`'s + docstring), so whichever event fires next must release everything + accumulated since the last release, not just its own hop's key. + + Hop 3's eventual success is fired as a child task of the same admission + chain -- exactly like litellm's real dispatch, where `wrapper_async` + create_task's the success path and `LoggingWorker.enqueue` explicitly + propagates the calling context -- to prove the fix survives the actual + task boundary a real success event crosses in production, not just a + same-coroutine call that would pass regardless of whether the pending + keys lived on a real shared object or an ordinary per-task variable. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=2) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + async def _one_logical_request(): + # All three hops of this one logical request share the same + # model_call_details, exactly as real fallback hops share one + # Logging object -- only litellm_call_id differs per hop. + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + + # Hop 1 admits and fails; its failure event is the one that fires + # (dedup allows exactly the first failure through), releasing its + # own key immediately. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + kwargs["standard_logging_object"] = {"model_group": "grp"} + await limiter.async_log_failure_event( + kwargs=kwargs, + response_obj=None, + start_time=0, + end_time=0, + ) + + # Hop 2 (a retry or fallback) admits and also fails, but -- per + # litellm's dedup -- no async_log_failure_event call follows it. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + + # Hop 3 admits and succeeds. Its success event, dispatched as a + # child task (mirroring the real worker hop), must release both + # hop 2's still-pending reservation and its own. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + + async def _hop_3_success_event(): + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 0, + "response_cost": 0, + } + await limiter.async_log_success_event( + kwargs=kwargs, + response_obj=None, + start_time=0, + end_time=0, + ) + + await asyncio.create_task(_hop_3_success_event()) + + await asyncio.create_task(_one_logical_request()) + await asyncio.sleep(0) + + # Full capacity (2) is free again -- both hop 2's leaked reservation and + # hop 3's own were released. If the earlier hop's leaked reservation + # hadn't been released too, only one of these two admissions would succeed. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_next_hops_admission_releases_a_prior_hops_leaked_reservation(time_controller): + """ + Regression test for a leak that a success/failure-event-only release + strategy can never close: litellm's has_logged_async_failure dedup lets + exactly one hop's async_log_failure_event fire per logical request (see + test_concurrency_released_for_every_hop_across_a_real_task_boundary), so + a hop that fails *after* that one event has already fired gets no + failure event of its own at all -- not "delayed until the next event", + genuinely never. Only the next hop's own admission call is guaranteed to + run afterward, so release must happen there, not wait for some later + success/failure event that this specific hop will never get. + + Concurrency limit of 1 makes this observable directly: if hop 2's + admission doesn't release hop 1's leaked reservation before checking its + own, it raises ProxyRateLimitError against a bucket that's actually free. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + + # Hop 1 admits (the only slot) and then fails with no failure event ever + # following it -- simulating every hop after litellm's one dedup-allowed + # failure event has already fired for an earlier hop of this request. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + # Hop 2's own admission call must release hop 1's stale reservation + # before checking its own -- if it didn't, this raises ProxyRateLimitError. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_controller): + """ + Veria AI finding: Router.abatch_completion's comma-separated multi-model + dispatch runs each model concurrently as its own asyncio.Task, but every + branch is handed the identical litellm_logging_obj (the proxy attaches + one to the request before the comma-split), so two genuinely concurrent + branches share one model_call_details. Before the fix, a second branch's + admission-time stale-reservation cleanup couldn't tell that apart from + an earlier, already-failed hop of its own retry chain, so it released + the first branch's still-live reservation and let both branches through + a concurrency limit of 1. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + first_admitted = asyncio.Event() + + async def _branch_one() -> None: + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + first_admitted.set() + # Held "in flight" while the second branch's admission runs, exactly + # like two concurrently in-flight provider calls. + await asyncio.sleep(0.05) + + async def _branch_two() -> None: + await first_admitted.wait() + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + results = await asyncio.gather( + asyncio.create_task(_branch_one()), asyncio.create_task(_branch_two()), return_exceptions=True + ) + rejections = [result for result in results if isinstance(result, ProxyRateLimitError)] + assert len(rejections) == 1 + + +def _request_limit_router(limit: int) -> "litellm.Router": + return litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + {"name": "per_period", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300} + ] + } + }, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_next_hops_admission_refunds_a_prior_failed_hops_request_increment(time_controller): + """ + Regression test for Cursor Bugbot's "fallback hops burn request budget" + finding on PR #36541, live-confirmed against a real proxy: a "requests" + limit is meant to cap logical client requests, not internal routing + attempts, but without a refund a chain that fails once before succeeding + burned 2 units of a 1-request-per-period budget for one logical call -- + live reproduction showed the retry's own admission rejected with + current=1.0 limit=1.0 even though the client only made one call. + + Concurrency's next-hop-releases-the-prior-hop's-stale-reservation pattern + (see test_next_hops_admission_releases_a_prior_hops_leaked_reservation) + generalizes cleanly here: since Router only re-enters admission for a hop + that already failed, the prior hop's own "requests" increment must be + refunded there too, before this hop's own check runs. + """ + limiter = _make_limiter(time_controller) + router = _request_limit_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + + # Hop 1 admits (the only unit) then fails -- no failure event follows, + # mirroring the "already consumed litellm's one dedup-allowed failure + # event" scenario the sibling concurrency test documents. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + # Hop 2's own admission must refund hop 1's now-stale "requests" + # increment before checking its own -- if it didn't, this raises + # ProxyRateLimitError against a bucket a real client only asked to use + # once. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_next_hops_admission_refunds_a_request_increment_even_after_the_first_hops_own_failure_event_fires( + time_controller, +): + """ + Tighter regression than the test above: this reproduces the exact live + failure this fix first shipped with. litellm's has_logged_async_failure + dedup allows exactly the *first* failing hop's own async_log_failure_event + through -- unlike a hop after that one, hop 1 here genuinely gets a real + failure event, not silence. An earlier version of this fix popped + _PENDING_REQUEST_INCREMENTS_FIELD in async_log_failure_event "for + hygiene", discarding hop 1's entry before hop 2's own admission + (_release_stale_hop_reservations) ever got a chance to refund it -- + silently and permanently stranding the charge, so hop 2 was rejected + against a bucket a real client only asked to use once, live-confirmed + against a real proxy. async_log_failure_event must leave this field + completely untouched. + """ + limiter = _make_limiter(time_controller) + router = _request_limit_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + # Hop 1's own, real failure event -- the one has_logged_async_failure + # lets through. + await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # Hop 2's own admission must still refund hop 1's now-stale "requests" + # increment before checking its own. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_successful_hops_own_request_increment_is_not_refunded(time_controller): + """ + The fix above must not swing the other way and refund every hop's + "requests" increment unconditionally -- exactly one unit must survive + per logical request, or the limit stops limiting anything. Simulates the + full lifecycle (admission, then the success event a real request would + fire) and confirms a second, unrelated logical request against the same + tag is correctly rejected: the first request's own successful hop + already spent the only unit for this period. + """ + limiter = _make_limiter(time_controller) + router = _request_limit_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + fresh_request_kwargs, _fresh_kwargs = _call_context(["end_user_id:u1"]) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=fresh_request_kwargs + ) + + +@pytest.mark.asyncio +async def test_a_hops_own_rejection_on_a_different_check_does_not_undercount_the_prior_hops_request_charge( + time_controller, +): + """ + Regression test for Cursor Bugbot's follow-up finding on this exact fix: + an earlier version refunded the prior hop's "requests" charge + unconditionally at the top of the next hop's admission, before knowing + whether that next hop would itself be admitted. If the next hop then + failed a *different* check (here, concurrency) before ever reaching its + own requests renewal, the refund had already committed with nothing to + replace it -- a logical request that genuinely made one real attempt + (hop 1) would end up charged zero, letting a caller bypass the requests + cap simply by having a later hop collide with someone else's + concurrency slot. + + Fixed by folding the renewal into the same all-or-nothing atomic batch + as every other check on that hop: a "requests" key matching an earlier + hop's charge renews at zero net cost instead of being refunded first, + so a batch-wide rollback (concurrency's own rejection here) refunds that + zero-cost renewal -- a genuine no-op -- leaving hop 1's real charge + exactly as it was. + """ + # Two independent tag identities: "requests" is scoped to end_user_id + # (private to our own request, never shared with the unrelated + # contender below), "concurrency" is scoped to a separate shared_pool + # tag that both our request and the unrelated contender carry, so they + # compete for the same slot without also colliding on the requests cap. + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_period", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + }, + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "shared_pool", "limit": 1, "period_seconds": 300}] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # Hop 1 of our request admits (claiming both the requests unit and the + # only concurrency slot), then fails for real -- its concurrency + # reservation is released the normal way, but its requests charge is + # left queued as this-hop's-own-charge, not refunded. + request_kwargs, kwargs = _call_context(["end_user_id:u1", "shared_pool:pool-a"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # A second, unrelated request -- no end_user_id tag at all, so it never + # touches the requests bucket -- now claims the concurrency slot our + # hop 1 just released, and holds it. + other_request_kwargs, other_kwargs = _call_context(["shared_pool:pool-a"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=other_request_kwargs + ) + + # Hop 2 of our original request: its own "requests" renewal would + # trivially succeed alone (net zero cost), but the concurrency slot is + # now held by the unrelated request above, so the whole atomic batch + # must reject -- and must NOT leave hop 1's requests charge refunded. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # The unrelated request finishes, freeing the concurrency slot again. + await limiter.async_log_success_event(kwargs=other_kwargs, response_obj=None, start_time=0, end_time=0) + + # A fresh probe against the same end_user_id tag, with concurrency now + # free, must still be rejected by the requests cap: hop 1's real attempt + # already spent the only unit for this period, and it must not have + # been silently erased by hop 2's unrelated, different-check rejection. + probe_request_kwargs, _probe_kwargs = _call_context(["end_user_id:u1"]) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=probe_request_kwargs + ) + + +@pytest.mark.asyncio +async def test_a_third_hops_admission_still_recognizes_a_charge_that_survived_a_middle_hops_rejection( + time_controller, +): + """ + Regression test for Cursor Bugbot's follow-up finding on the fix above: + an earlier version had _release_stale_hop_reservations *pop* the pending + "requests" entry, re-queuing it only after this hop's own atomic batch + fully succeeded. A middle hop that failed *before* reaching that point + (exactly the concurrency-rejection scenario the test above covers) left + the real counter correctly charged but the bookkeeping field empty, so + a *third* hop's own peek came back with nothing to renew and charged a + fresh unit on top of the one still sitting in the real counter -- + doubling the charge (or, with a tighter limit, a false 429) despite the + fix that was supposed to prevent exactly that. + + Reuses the same request_kwargs (the same model_call_details) across all + three hops -- unlike the probe above, which deliberately uses a fresh, + independent context and so can't distinguish "the real counter is + correct" from "the bookkeeping that lets a future hop recognize it is + intact"; only a third hop sharing the same chain's own bookkeeping can. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_period", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + }, + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "shared_pool", "limit": 1, "period_seconds": 300}] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # Hop 1 admits (claiming the requests unit and the only concurrency + # slot), then fails for real. + request_kwargs, kwargs = _call_context(["end_user_id:u1", "shared_pool:pool-a"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # An unrelated request claims the now-free concurrency slot and holds it. + other_request_kwargs, other_kwargs = _call_context(["shared_pool:pool-a"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=other_request_kwargs + ) + + # Hop 2 of our original request rejects on concurrency (not requests) -- + # its own admission never reaches its own successful queuing step. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # The unrelated request finishes, freeing the concurrency slot. Release + # runs as a background task on success (see async_log_success_event's + # own implementation), so let it actually complete before hop 3 checks. + await limiter.async_log_success_event(kwargs=other_kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # Hop 3 of our original request, same request_kwargs: if the bookkeeping + # survived hop 2's rejection, this renews at zero cost and succeeds. If + # it was lost, this charges a fresh unit on top of hop 1's still-live + # real charge and wrongly rejects. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_own_rejection_does_not_release_a_live_reservation(time_controller): + """ + Security regression test: Router.async_callback_filter_deployments fires + async_log_failure_event for an exception raised from inside + async_filter_deployments itself (its own except block calls + logging_obj.async_failure_handler before re-raising) -- not only for an + actual provider-call failure. A rejection this hook raises for being + over its own limit never reserved anything for that specific attempt + (_atomic_check_and_increment already refunds any of its own earlier + admissions synchronously whenever it rejects), so releasing anyway would + decrement a live reservation belonging to a different, genuinely + in-flight request sharing the same tag -- letting a caller free up + capacity simply by retrying against an already-full bucket, no + coordination with another request required. + + The holder and the rejected attempt are modeled as two separate tasks + (matching how two independent real requests are always isolated in + production, each in its own asyncio task) so this actually exercises the + explicit ProxyRateLimitError guard rather than the ContextVar's own + per-task isolation, which would otherwise mask the same bug: two + admissions made directly in one shared coroutine would (correctly, but + for the wrong reason) never be able to explain away the bug this test is + for. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # One legitimate request, in its own task, holds the only slot. + async def _admit(): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + await asyncio.create_task(_admit()) + + # A second, unrelated request (its own task) is rejected, and Router's + # own exception handling fires async_log_failure_event for it, exactly + # as Router.async_callback_filter_deployments does. + async def _reject_and_fire_failure_event(): + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + await limiter.async_log_failure_event( + kwargs={ + "exception": exc_info.value, + "standard_logging_object": {"model_group": "grp"}, + "metadata": {"tags": ["end_user_id:u1"]}, + }, + response_obj=None, + start_time=0, + end_time=0, + ) + + await asyncio.create_task(_reject_and_fire_failure_event()) + + # The first request's reservation must still be held: a third attempt is + # still rejected. If the rejection's failure event had wrongly released + # it, this would wrongly admit instead. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + +# --------------------------------------------------------------------------- +# tokens / dollars -- read-then-account-on-success rejection path +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_token_limit_rejects_once_bucket_is_seeded_at_limit(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + now = time_controller.now().timestamp() + key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=1000) + await limiter.internal_usage_cache.async_set_cache(key=key, value=1000, ttl=86400, litellm_parent_otel_span=None) + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert exc_info.value.detail["type"] == "tokens" + assert exc_info.value.detail["limit_name"] == "daily" + + +@pytest.mark.asyncio +async def test_dollar_limit_rejects_once_bucket_is_seeded_at_limit(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "dollar_limits": { + "limits": [{"name": "monthly", "tag_id": "team_id", "limit": 50.0, "period_seconds": 2592000}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + now = time_controller.now().timestamp() + key = _expected_bucket_key("grp", "dollars", "monthly", "team_id", "t1", 2592000, now, limit=50.0) + await limiter.internal_usage_cache.async_set_cache(key=key, value=50.0, ttl=2592000, litellm_parent_otel_span=None) + + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["team_id:t1"]}}, + ) + assert exc_info.value.detail["type"] == "dollars" + assert exc_info.value.detail["tag_value"] == "t1" + + +# --------------------------------------------------------------------------- +# Real Redis -- the atomic Lua path, not just the in-memory fallback +# --------------------------------------------------------------------------- + + +def _redis_limiter(time_controller: TimeController): + import os + + from litellm.caching.redis_cache import RedisCache + + redis_host = os.getenv("REDIS_HOST") + redis_port = os.getenv("REDIS_PORT") + if not redis_host or not redis_port: + pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set") + redis_cache = RedisCache(host=redis_host, port=int(redis_port), password=os.getenv("REDIS_PASSWORD")) + dual_cache = DualCache(redis_cache=redis_cache) + return _PROXY_ModelBasedTagRateLimitsHook( + internal_usage_cache=dual_cache, time_provider=time_controller.now + ), redis_cache + + +@pytest.mark.asyncio +async def test_redis_backed_requests_admission_is_race_free_under_genuine_concurrency(time_controller): + """ + Same race-freedom guarantee as the in-memory test, but against a real + Redis instance so the Lua script path (not just the asyncio.Lock + fallback) is exercised -- this is the code path every multi-instance + proxy deployment actually runs. + """ + limiter, redis_cache = _redis_limiter(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 5, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + tag = f"redis-race-{uuid.uuid4().hex}" + + async def attempt(): + try: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:{tag}"]}}, + ) + return True + except ProxyRateLimitError: + return False + + results = await asyncio.gather(*(attempt() for _ in range(20))) + assert sum(results) == 5 + + +@pytest.mark.asyncio +async def test_redis_backed_cross_unit_rejection_does_not_leave_a_phantom_increment(time_controller): + """Redis-Lua-script equivalent of the in-memory phantom-increment regression test.""" + limiter, redis_cache = _redis_limiter(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 10, "period_seconds": 60}] + }, + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + tag = f"redis-phantom-check-{uuid.uuid4().hex}" + + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:{tag}"]}}, + ) + with pytest.raises(ProxyRateLimitError) as exc_info: + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:{tag}"]}}, + ) + assert exc_info.value.detail["type"] == "concurrency" + + now = time_controller.now().timestamp() + request_key = _expected_bucket_key("grp", "requests", "per_minute", "end_user_id", tag, 60, now, limit=10) + requests_value = await limiter.internal_usage_cache.async_get_cache(key=request_key, litellm_parent_otel_span=None) + assert (float(requests_value) if requests_value is not None else 0.0) == 1.0 + + # cleanup: this key persists in the shared scratch Redis instance beyond the test's TTL otherwise + await redis_cache.async_delete_cache(key=request_key) + + +@pytest.mark.asyncio +async def test_redis_backed_token_admission_sees_increments_the_in_memory_cache_missed(time_controller): + """ + Success accounting increments a token bucket straight through a Lua + script on redis_cache, bypassing DualCache/InternalUsageCache entirely -- + that write never touches the in-memory layer. Once an earlier read has + backfilled that same key into the in-memory cache, DualCache's own + async_batch_get_cache treats that non-None in-memory hit as authoritative + and never re-checks Redis, so every later admission would see the same + frozen snapshot while the real Redis counter keeps climbing underneath + it, silently admitting traffic well past the configured token limit. + """ + limiter, redis_cache = _redis_limiter(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 100, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + tag = f"redis-stale-check-{uuid.uuid4().hex}" + request_kwargs = {"metadata": {"tags": [f"end_user_id:{tag}"]}} + + async def _charge(tokens: float) -> None: + await limiter.async_log_success_event( + kwargs={ + "metadata": {"tags": [f"end_user_id:{tag}"]}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": tokens, + "response_cost": 0, + }, + }, + response_obj=None, + start_time=0, + end_time=0, + ) + # The actual Redis increment is dispatched as a background task (see + # _BACKGROUND_TASKS), so it needs a beat to actually run. + await asyncio.sleep(0.05) + + # First admission: bucket doesn't exist in Redis yet, so this read finds + # nothing to backfill into the in-memory cache either. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + await _charge(90) + + # Second admission: this read is the one that backfills the in-memory + # cache with the real (90) value read from Redis. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + await _charge(90) # real Redis total is now 180, well past the limit of 100 + + # Third admission must see the real (180) total and reject -- not the + # frozen 90 the in-memory cache captured on the previous read. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "per_minute", "end_user_id", tag, 60, now) + await redis_cache.async_delete_cache(key=token_key) + + +@pytest.mark.asyncio +async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + Bugbot finding: TAG_RL_CHECK_AND_INCR_SCRIPT only ran EXPIRE when a key + had no TTL at all, so a concurrency counter's expiry was fixed from its + first admission and never pushed out by later ones. A concurrency bucket + isn't epoch-windowed like requests/tokens/dollars -- its TTL exists only + as a crash-safety net for a reservation whose explicit release never + runs -- so a still-active bucket receiving continuous admissions must + keep extending that TTL, or it expires mid-flight under sustained + traffic, silently admitting past the cap. + """ + limiter, redis_cache = _redis_limiter(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + key = f"{{tag_rl:test:ttl-refresh:{uuid.uuid4().hex}}}:inflight" + cache = limiter.internal_usage_cache + try: + # A short, fixed ttl (bypassing _ttl_for's 3600s safety floor, which + # would make a real-time before/after comparison too slow to assert + # on deterministically) with refresh_ttl=True, matching how a + # concurrency check is actually admitted. + admitted, _ = await limiter._check_and_increment_one( + cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True + ) + assert admitted + ttl_after_first_admission = await redis_cache.redis_async_client.ttl(key) + assert ttl_after_first_admission > 0 + + await asyncio.sleep(2) + + # A second admission on the same still-live key, most of the way + # through the first admission's ttl, must push the ttl back out to + # the full window again, not leave it counting down toward zero. + admitted, _ = await limiter._check_and_increment_one( + cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True + ) + assert admitted + ttl_after_second_admission = await redis_cache.redis_async_client.ttl(key) + assert ttl_after_second_admission >= 2 + finally: + await redis_cache.async_delete_cache(key=key) + + +@pytest.mark.asyncio +async def test_in_memory_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + Bugbot finding: the Redis path's refresh_ttl fix above was never mirrored + onto the in-memory fallback, which called async_set_cache unconditionally + -- InMemoryCache.allow_ttl_override leaves a still-live ttl untouched, so + a concurrency counter's expiry stayed fixed from its first admission even + with refresh_ttl=True, the same silent-past-the-cap failure mode the + Redis fix closed. + """ + limiter = _make_limiter(time_controller) + cache = limiter.internal_usage_cache + in_memory_cache = cache.dual_cache.in_memory_cache + key = f"tag_rl:test:in-memory-ttl-refresh:{uuid.uuid4().hex}" + + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = in_memory_cache.ttl_dict[key] + + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = in_memory_cache.ttl_dict[key] + + assert ttl_after_second_admission > ttl_after_first_admission + + +# --------------------------------------------------------------------------- +# team_public_model_name alias -- index lookup must not miss +# --------------------------------------------------------------------------- + + +def test_build_limits_index_is_also_keyed_by_team_public_model_name(): + """ + Router threads a team's public alias, not the deployment's own + model_name, into async_filter_deployments's `model` param when a caller + requests via that alias (Router never rewrites it for this path, unlike + model_group_alias). The index must resolve either name to the same + configured limits, or a team-aliased chain's limits are silently never + checked. + + Security regression: Router.should_include_deployment also lets a + same-team (or team-unconstrained) caller reach this deployment by its + own internal model_name, not only the alias. Both paths must resolve to + the identical team_scope, or a caller could split its usage across two + independent buckets just by alternating which name it calls with. + """ + deployment = _deployment( + "real-model-name", + "dep-1", + {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}}, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + index = _build_limits_index([deployment]) + by_name = index.resolve("real-model-name", team_id=None) + by_alias = index.resolve("team-alias-name", team_id="team-1") + assert by_name != () + assert [c.entry for c in by_name] == [c.entry for c in by_alias] + assert by_name[0].team_scope == "team-1" + assert by_alias[0].team_scope == "team-1" + + +def test_build_limits_index_computes_identical_bucket_key_for_alias_and_internal_model_name(): + """ + Bugbot finding: matching team_scope alone does not unify the two paths' + buckets, because _hash_tag hashes the caller-visible model_group by + default, and that name is `real-model-name` on the by_model_name path but + `team-alias-name` on the by_team_alias path. Both entries must also share + the identical resolved_group (the team's own alias) so _hash_tag hashes + both under one name -- otherwise a team calling its own internal + model_name lands on a different Redis counter than the same team calling + its public alias, doubling its effective quota. + """ + deployment = _deployment( + "real-model-name", + "dep-1", + {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}}, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + index = _build_limits_index([deployment]) + by_name = index.resolve("real-model-name", team_id=None)[0] + by_alias = index.resolve("team-alias-name", team_id="team-1")[0] + + key_via_internal_name = _bucket_key("real-model-name", by_name, tag_value="u1", bucket_id=0) + key_via_alias = _bucket_key("team-alias-name", by_alias, tag_value="u1", bucket_id=0) + assert key_via_internal_name == key_via_alias + + +def test_resolve_any_preserves_team_alias_resolved_group_through_routing_group_fallback(): + """ + Cursor Bugbot finding: resolve_any's routing-group fallback always + restamped resolved_group with the candidate's own model_name, discarding + the team_public_model_name _build_limits_index already stamped there -- + so a team-owned deployment reachable both via its team alias and via a + routing group split one team's usage across two buckets depending on + which path a caller took. + """ + deployment = _deployment( + "real-model-name", + "dep-1", + {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}}, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + index = _build_limits_index([deployment]) + + via_alias = index.resolve("team-alias-name", team_id="team-1")[0] + via_routing_group_fallback = index.resolve_any( + "my-group", team_id="team-1", candidate_model_names=["real-model-name"] + )[0] + + key_via_alias = _bucket_key("team-alias-name", via_alias, tag_value="u1", bucket_id=0) + key_via_fallback = _bucket_key("my-group", via_routing_group_fallback, tag_value="u1", bucket_id=0) + assert key_via_alias == key_via_fallback + + +def test_build_limits_index_preserves_key_ttl_seconds_and_max_in_memory_cache_size(): + """ + Regression test: _configured_limit_for_signature used to reconstruct a + fresh TagRateLimitEntry from a 5-field dedup signature that didn't + include key_ttl_seconds or max_in_memory_cache_size, silently resetting + both to None for every entry that went through the real indexing path + (which is every entry reachable from async_filter_deployments / + async_log_success_event) -- only entries built directly in a test, never + through _build_limits_index, kept their configured values. + """ + deployment = _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "user_cap", + "tag_id": "end_user_id", + "limit": 5, + "period_seconds": 60, + "key_ttl_seconds": 120, + "max_in_memory_cache_size": 500, + } + ] + } + }, + ) + index = _build_limits_index([deployment]) + configured = index.resolve("grp", team_id=None) + assert len(configured) == 1 + assert configured[0].entry.key_ttl_seconds == 120 + assert configured[0].entry.max_in_memory_cache_size == 500 + + +def test_build_limits_index_treats_a_duplicated_entry_on_one_deployment_as_chain_wide(): + """ + Regression test: a single deployment declaring the identical + concurrency_limits entry twice (a config duplicate) used to append that + deployment's id twice, inflating len(declaring_ids) past + total_deployments. That made is_chain_wide false even though every + deployment (there's only one) actually agreed on the entry, and for + concurrency a non-chain-wide entry is silently dropped entirely -- + disabling enforcement rather than degrading it. + """ + deployment = _deployment( + "grp", + "dep-1", + { + "concurrency_limits": { + "limits": [ + {"name": "inflight", "tag_id": "end_user_id", "limit": 5, "period_seconds": 300}, + {"name": "inflight", "tag_id": "end_user_id", "limit": 5, "period_seconds": 300}, + ] + } + }, + ) + index = _build_limits_index([deployment]) + configured = index.resolve("grp", team_id=None) + assert len(configured) == 1 + assert configured[0].deployment_scope is None # chain-wide, not dropped + + +def test_build_limits_index_keeps_different_teams_same_alias_separate(): + """ + `team_public_model_name` is only unique per team: Router itself lets two + different teams publish the identical alias string for different + deployments, resolving each caller's own team's deployment by + `(team_id, name)` rather than by name alone. Keying the limits index by + name alone would let one team's config silently overwrite another's. + """ + team_a = _deployment( + "model-a", "dep-a", {"token_limits": {"limits": [{"name": "daily", "limit": 100, "period_seconds": 86400}]}} + ) + team_a["model_info"]["team_id"] = "team-a" + team_a["model_info"]["team_public_model_name"] = "shared-alias" + + team_b = _deployment( + "model-b", "dep-b", {"token_limits": {"limits": [{"name": "daily", "limit": 999, "period_seconds": 86400}]}} + ) + team_b["model_info"]["team_id"] = "team-b" + team_b["model_info"]["team_public_model_name"] = "shared-alias" + + index = _build_limits_index([team_a, team_b]) + resolved_a = index.resolve("shared-alias", team_id="team-a") + resolved_b = index.resolve("shared-alias", team_id="team-b") + assert resolved_a[0].entry.limit == 100 + assert resolved_b[0].entry.limit == 999 + + +def test_bucket_key_differs_across_teams_sharing_an_alias_and_identical_limit_config(): + """ + Two teams that happen to publish the identical team_public_model_name + AND configure an identically-named, identically-valued limit must not + land on the same Redis bucket -- team_public_model_name is only unique + per team, so this is a realistic collision, not a contrived one. + """ + team_a = _deployment( + "model-a", "dep-a", {"request_limits": {"limits": [{"name": "per_minute", "limit": 5, "period_seconds": 60}]}} + ) + team_a["model_info"]["team_id"] = "team-a" + team_a["model_info"]["team_public_model_name"] = "shared-alias" + + team_b = _deployment( + "model-b", "dep-b", {"request_limits": {"limits": [{"name": "per_minute", "limit": 5, "period_seconds": 60}]}} + ) + team_b["model_info"]["team_id"] = "team-b" + team_b["model_info"]["team_public_model_name"] = "shared-alias" + + index = _build_limits_index([team_a, team_b]) + limit_a = index.resolve("shared-alias", team_id="team-a")[0] + limit_b = index.resolve("shared-alias", team_id="team-b")[0] + assert limit_a.entry == limit_b.entry # identical configuration, by construction + + key_a = _bucket_key("shared-alias", limit_a, tag_value="same-caller-tag", bucket_id=0) + key_b = _bucket_key("shared-alias", limit_b, tag_value="same-caller-tag", bucket_id=0) + assert key_a != key_b + + inflight_a = _inflight_key("shared-alias", limit_a, tag_value="same-caller-tag") + inflight_b = _inflight_key("shared-alias", limit_b, tag_value="same-caller-tag") + assert inflight_a != inflight_b + + +def test_build_limits_index_merges_alias_limits_across_different_model_names(): + """ + litellm auto-generates each team-added deployment's own internal + model_name as model_name_{team_id}_{uuid}, so multiple deployments + sharing one team_public_model_name alias routinely have different + model_name values -- Router's own team_model_to_deployment_indices + aggregates them by (team_id, alias) regardless of that. Computing alias + limits once per model_name group and keying the alias to whichever + group happened to declare it would silently drop every other same-alias + group's limits: with two deployments under different model_names but + the same alias, only the entry declared by whichever model_name group + is processed last would survive. + """ + dep_a = _deployment( + "model_name_team1_aaa", + "dep-a", + {"token_limits": {"limits": [{"name": "daily", "limit": 100, "period_seconds": 86400}]}}, + ) + dep_a["model_info"]["team_id"] = "team-1" + dep_a["model_info"]["team_public_model_name"] = "shared-alias" + + dep_b = _deployment( + "model_name_team1_bbb", + "dep-b", + {"dollar_limits": {"limits": [{"name": "monthly", "limit": 50.0, "period_seconds": 2592000}]}}, + ) + dep_b["model_info"]["team_id"] = "team-1" + dep_b["model_info"]["team_public_model_name"] = "shared-alias" + + index = _build_limits_index([dep_a, dep_b]) + resolved = index.resolve("shared-alias", team_id="team-1") + units = {c.unit for c in resolved} + assert units == {"tokens", "dollars"} + + +@pytest.mark.asyncio +async def test_filter_deployments_enforces_limit_when_called_with_team_alias(time_controller): + limiter = _make_limiter(time_controller) + deployment = _deployment( + "real-model-name", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # Router passes the alias as `model`, not "real-model-name", and threads + # the caller's team_id through request metadata. + request_kwargs = {"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}} + await limiter.async_filter_deployments( + model="team-alias-name", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="team-alias-name", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_shares_quota_across_alias_and_internal_model_name(time_controller): + """ + Bugbot finding on a prior fix for this same drift: stamping the identical + team_scope on both paths was not enough, since _hash_tag still hashes the + caller-visible name (the alias here, the deployment's own internal + model_name there) by default. A team calling through its public alias and + the identical team calling through the deployment's own internal + model_name must draw from the same bucket, or the team gets one quota per + name it happens to call with -- doubling its real limit. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "real-model-name", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # First call exhausts the limit of 1 via the team's public alias. + await limiter.async_filter_deployments( + model="team-alias-name", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}}, + ) + # Router.should_include_deployment also lets the same team reach this + # deployment by its own internal model_name; that call must be rejected + # against the alias call's own bucket, not admitted into a fresh one. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="real-model-name", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}}, + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_does_not_cross_team_alias_boundary(time_controller): + """ + Two teams sharing the same team_public_model_name must not share a + counter: a caller on team-b hitting the alias must not be limited (or + counted) by team-a's configured limit and usage. + """ + limiter = _make_limiter(time_controller) + team_a = _deployment( + "model-a", + "dep-a", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + team_a["model_info"]["team_id"] = "team-a" + team_a["model_info"]["team_public_model_name"] = "shared-alias" + + team_b = _deployment( + "model-b", + "dep-b", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 5, "period_seconds": 86400}] + } + }, + ) + team_b["model_info"]["team_id"] = "team-b" + team_b["model_info"]["team_public_model_name"] = "shared-alias" + + router = litellm.Router(model_list=[team_a, team_b]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # team-a exhausts its own limit of 1. + await limiter.async_filter_deployments( + model="shared-alias", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-a"}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="shared-alias", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-a"}}, + ) + + # team-b, same alias string and same tag value, is unaffected by team-a's exhausted limit. + await limiter.async_filter_deployments( + model="shared-alias", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-b"}}, + ) + + +# --------------------------------------------------------------------------- +# concurrency_limits -- chain-wide only, divergent config dropped not leaked +# --------------------------------------------------------------------------- + + +def test_concurrency_divergent_config_is_dropped_not_scoped_per_deployment(): + """ + Unlike tokens/requests/dollars, a concurrency entry declared with + different values per deployment must not become a per-deployment-scoped + reservation -- that shape leaks (see the regression tests below this + plan superseded). It should be dropped entirely, with the group left + with no concurrency entry at all. + """ + deployments = [ + _deployment( + "grp", "dep-1", {"concurrency_limits": {"limits": [{"name": "inflight", "limit": 2, "period_seconds": 60}]}} + ), + _deployment( + "grp", "dep-2", {"concurrency_limits": {"limits": [{"name": "inflight", "limit": 5, "period_seconds": 60}]}} + ), + ] + configured = _build_group_limits(deployments, "concurrency") + assert configured == () + + +def test_concurrency_partial_declaration_is_dropped_not_scoped_per_deployment(): + deployments = [ + _deployment( + "grp", "dep-1", {"concurrency_limits": {"limits": [{"name": "inflight", "limit": 2, "period_seconds": 60}]}} + ), + _deployment("grp", "dep-2", {}), + ] + configured = _build_group_limits(deployments, "concurrency") + assert configured == () + + +def test_concurrency_identical_across_all_deployments_is_still_chain_wide(): + deployments = [ + _deployment( + "grp", "dep-1", {"concurrency_limits": {"limits": [{"name": "inflight", "limit": 2, "period_seconds": 60}]}} + ), + _deployment( + "grp", "dep-2", {"concurrency_limits": {"limits": [{"name": "inflight", "limit": 2, "period_seconds": 60}]}} + ), + ] + configured = _build_group_limits(deployments, "concurrency") + assert len(configured) == 1 + assert configured[0].deployment_scope is None + + +# --------------------------------------------------------------------------- +# concurrency TTL floor -- a short period_seconds must not shorten the +# self-heal safety TTL below the floor +# --------------------------------------------------------------------------- + + +def test_concurrency_ttl_floor_overrides_a_too_short_period_seconds(): + entry = TagRateLimitEntry(name="inflight", tag_id="end_user_id", limit=1, period_seconds=5) + configured_limit = _ConfiguredLimit(unit="concurrency", entry=entry, deployment_scope=None) + assert _PROXY_ModelBasedTagRateLimitsHook._ttl_for(configured_limit) == _CONCURRENCY_MIN_SAFETY_TTL_SECONDS + + +def test_concurrency_ttl_floor_does_not_shorten_a_longer_period_seconds(): + entry = TagRateLimitEntry( + name="inflight", tag_id="end_user_id", limit=1, period_seconds=_CONCURRENCY_MIN_SAFETY_TTL_SECONDS + 100 + ) + configured_limit = _ConfiguredLimit(unit="concurrency", entry=entry, deployment_scope=None) + assert _PROXY_ModelBasedTagRateLimitsHook._ttl_for(configured_limit) == _CONCURRENCY_MIN_SAFETY_TTL_SECONDS + 100 + + +# --------------------------------------------------------------------------- +# pending-concurrency-key field on model_call_details must survive a detached +# asyncio.create_task fork (e.g. litellm's own failure-logging dispatch), +# and a release must never sweep up a key a still-live sibling hop appended +# in the meantime. This dict-on-a-shared-object design is what replaced a +# contextvars.ContextVar-based holder that silently failed to release +# anything once release ran in a task that wasn't a descendant of admission's +# own task -- exactly what happens in the real proxy request pipeline (see +# _PENDING_CONCURRENCY_KEYS_FIELD's docstring). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_release_in_a_forked_task_is_visible_to_the_parent_context(time_controller): + limiter = _make_limiter(time_controller) + # Entries are (key, partition_key, admission_token) triples in production + # (see _queue_pending_reservations); the token is irrelevant to this + # specific release path (only_own_lineage defaults False here). + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} + + async def detached_release(): + return await limiter._pop_pending_concurrency_keys(model_call_details) + + released = await asyncio.create_task(detached_release()) + assert released == (("key1", None),) + + # The parent's own view of the same dict must see the release too. + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [] + + +@pytest.mark.asyncio +async def test_release_does_not_sweep_up_a_key_appended_after_its_snapshot(time_controller): + limiter = _make_limiter(time_controller) + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} + + async def detached_release_then_sibling_admits(): + released = await limiter._pop_pending_concurrency_keys(model_call_details) + # A sibling hop's admission, appending to the same shared dict, + # interleaved right after this release's snapshot was taken. + model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD].append(("key2", None, None)) + return released + + released = await asyncio.create_task(detached_release_then_sibling_admits()) + assert released == (("key1", None),) + # key2 must still be pending for its own hop's eventual release. + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("key2", None, None)] + + +@pytest.mark.asyncio +async def test_release_is_not_repeated_for_the_same_snapshot(time_controller): + limiter = _make_limiter(time_controller) + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} + first = await limiter._pop_pending_concurrency_keys(model_call_details) + second = await limiter._pop_pending_concurrency_keys(model_call_details) + assert first == (("key1", None),) + assert second == () + + +@pytest.mark.asyncio +async def test_release_only_own_lineage_leaves_a_concurrent_siblings_reservation_alone(time_controller): + """ + Bugbot/Veria AI finding: Router.abatch_completion's comma-separated + multi-model dispatch runs each model concurrently, each its own asyncio + Task, but every branch shares one litellm_logging_obj (the proxy attaches + it to the request before the comma-split), so a still-live sibling + branch's own reservation can sit in the same model_call_details. + only_own_lineage must release this context's own entry while leaving a + differently-lineaged (sibling branch's) entry untouched. + """ + limiter = _make_limiter(time_controller) + own_token = _current_admission_token() + sibling_token = object() + model_call_details: dict = { + _PENDING_CONCURRENCY_KEYS_FIELD: [("own-key", None, own_token), ("sibling-key", None, sibling_token)] + } + + released = await limiter._pop_pending_concurrency_keys(model_call_details, only_own_lineage=True) + assert released == (("own-key", None),) + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_token)] + + +# --------------------------------------------------------------------------- +# refund-on-rollback across differently-hash-tagged keys (Redis Cluster safety) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_cross_unit_refund_leaves_no_phantom_increment_in_memory(time_controller): + """ + In-memory equivalent of the Redis Cluster cross-slot fix: requests and + concurrency keys carry different hash tags by construction, so the + all-or-nothing guarantee across them must come from a refund, not a + single multi-key atomic call. Confirms the refund path itself (not just + the end observable behavior already covered by + test_cross_unit_rejection_does_not_leave_a_phantom_increment). + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 10, "period_seconds": 60}] + }, + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:refund-check"]}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:refund-check"]}}, + ) + + now = time_controller.now().timestamp() + request_key = _expected_bucket_key( + "grp", "requests", "per_minute", "end_user_id", "refund-check", 60, now, limit=10 + ) + value = await limiter.internal_usage_cache.async_get_cache(key=request_key, litellm_parent_otel_span=None) + assert (float(value) if value is not None else 0.0) == 1.0 + + +# --------------------------------------------------------------------------- +# release floors at zero -- never goes negative +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_release_floors_at_zero_instead_of_going_negative(time_controller): + limiter = _make_limiter(time_controller) + key = "{tag_rl:test:concurrency:floor:chain:u1}:inflight" + await limiter._decrement_floor_zero(limiter.internal_usage_cache, key, -1.0) + value = await limiter.internal_usage_cache.async_get_cache(key=key, litellm_parent_otel_span=None) + assert (float(value) if value is not None else 0.0) == 0.0 + + +# --------------------------------------------------------------------------- +# a failed refund must not block refunding the rest of the batch or raise +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_refund_failure_on_one_key_does_not_block_others_or_raise(time_controller): + """ + If `_decrement_floor_zero` fails for one key mid-rollback (e.g. a + transient Redis error), the failure must be logged and swallowed, not + raised: otherwise it would surface as an unhandled exception in place of + the clean rejection the caller expects, and would abort the loop before + refunding every other already-committed key in the same batch. + """ + failing_key = "{tag_rl:test:refund-fail:a}:requests" + other_key = "{tag_rl:test:refund-fail:b}:requests" + rejecting_key = "{tag_rl:test:refund-fail:c}:requests" + + class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook): + async def _decrement_floor_zero(self, cache, key: str, delta: float) -> None: + if key == failing_key: + raise RuntimeError("simulated transient redis failure") + await super()._decrement_floor_zero(cache, key, delta) + + flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now) + + failing_index, values = await flaky._atomic_check_and_increment( + [ + (flaky.internal_usage_cache, failing_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, other_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, rejecting_key, 0.0, 1.0, 60, False), + ] + ) + + assert failing_index == 2 + + other_value = await flaky.internal_usage_cache.async_get_cache(key=other_key, litellm_parent_otel_span=None) + assert (float(other_value) if other_value is not None else 0.0) == 0.0 + + +@pytest.mark.asyncio +async def test_exception_mid_batch_refunds_every_earlier_admission_before_propagating(time_controller): + """ + Regression test: a later key's own admission raising (a transient Redis + error, or this coroutine being cancelled mid-call) used to skip the + refund loop entirely, since it only ran on a normal rejection return. + An earlier admission in the same batch would then stay permanently + charged -- for concurrency, a leaked reservation the caller never gets + to release, incorrectly throttling that tag until the 1-hour safety TTL + expires. + """ + admitted_key = "{tag_rl:test:exception-refund:a}:requests" + raising_key = "{tag_rl:test:exception-refund:b}:requests" + + class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook): + async def _check_and_increment_one( + self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool + ): + if key == raising_key: + raise RuntimeError("simulated transient redis failure") + return await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) + + flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now) + + with pytest.raises(RuntimeError): + await flaky._atomic_check_and_increment( + [ + (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60, False), + ] + ) + + admitted_value = await flaky.internal_usage_cache.async_get_cache(key=admitted_key, litellm_parent_otel_span=None) + assert (float(admitted_value) if admitted_value is not None else 0.0) == 0.0 + + +@pytest.mark.asyncio +async def test_a_raising_keys_own_ambiguous_outcome_is_never_refunded(time_controller): + """ + Regression test for a bug this exact fix briefly introduced: a key can + commit its own increment (e.g. Redis runs the INCRBY) and still have + the call raise if the response back to us is lost, so a raise never + proves that key's own attempt didn't commit. But these are shared, + chain-wide buckets with no per-request ownership tracking, so + decrementing on that guess is just as likely to erase a *different*, + legitimately-admitted concurrent request's charge on the same key as it + is to undo our own -- an attacker could repeatedly cancel requests to + erase other callers' charges and exceed the configured limit. The + raising key's own outcome must never be refunded, only strictly earlier + (confirmed-safe) admissions in the same batch. + """ + admitted_key = "{tag_rl:test:ambiguous-no-refund:a}:requests" + raising_key = "{tag_rl:test:ambiguous-no-refund:b}:requests" + + class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook): + async def _check_and_increment_one( + self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool + ): + if key == raising_key: + # Simulate Redis committing the increment before the + # response is lost: the write actually happens... + await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) + # ...but the caller never finds out. + raise RuntimeError("simulated lost response after a committed redis write") + return await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) + + flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now) + + with pytest.raises(RuntimeError): + await flaky._atomic_check_and_increment( + [ + (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60, False), + ] + ) + + # The earlier, confirmed-successful admission in this same batch is + # always safe to refund. + admitted_value = await flaky.internal_usage_cache.async_get_cache(key=admitted_key, litellm_parent_otel_span=None) + assert (float(admitted_value) if admitted_value is not None else 0.0) == 0.0 + + # The raising key's own committed increment must survive -- refunding + # it would be indistinguishable from erasing a different request's + # legitimate charge on the same shared bucket. + raising_key_value = await flaky.internal_usage_cache.async_get_cache(key=raising_key, litellm_parent_otel_span=None) + assert float(raising_key_value) == 1.0 + + +# --------------------------------------------------------------------------- +# scope_by_key_hash -- opt-in per-calling-key bucket separation +# --------------------------------------------------------------------------- + + +def _concurrency_router_scoped_by_key(limit: int) -> "litellm.Router": + return litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "concurrency_limits": { + "limits": [ + { + "name": "inflight", + "tag_id": "end_user_id", + "limit": limit, + "period_seconds": 300, + "scope_by_key_hash": True, + } + ] + } + }, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_request_limit_scope_by_key_hash_gives_independent_counters_per_key(time_controller): + """ + scope_by_key_hash=True: the identical tag value sent by two different + calling keys must get independent request counters -- exhausting keyA's + limit must not affect keyB's admission for the same tag. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "per_minute", + "tag_id": "end_user_id", + "limit": 2, + "period_seconds": 60, + "scope_by_key_hash": True, + } + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + for _ in range(2): + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + + # keyB, identical tag value, is unaffected -- it gets its own bucket and + # can admit up to the same limit independently of keyA's exhausted one. + for _ in range(2): + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}}, + ) + + +@pytest.mark.asyncio +async def test_scope_by_key_hash_composes_with_max_in_memory_cache_size_and_key_ttl_seconds_overrides(time_controller): + """ + scope_by_key_hash must keep working when combined with the two other + per-entry overrides, going through the real _build_limits_index path + (not a hand-built _ConfiguredLimit) -- this is exactly the path the + signature-reconstruction bug silently broke key_ttl_seconds and + max_in_memory_cache_size on, so it's worth covering in combination + rather than trusting the fields compose correctly in isolation. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "per_minute", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "scope_by_key_hash": True, + "max_in_memory_cache_size": 5, + "key_ttl_seconds": 120, + } + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + configured = limiter._index.get(router).resolve("grp", team_id=None) + assert configured[0].entry.max_in_memory_cache_size == 5 + assert configured[0].entry.key_ttl_seconds == 120 + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + assert result == healthy + + # keyA is now at its per-key limit of 1. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + + # keyB, identical tag value, still gets its own independent bucket on + # the same (overridden) cache partition. + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_request_limit_without_scope_by_key_hash_still_shares_one_counter(time_controller): + """ + Regression guard: scope_by_key_hash defaults to False, so today's + existing behavior -- the bucket is shared across every key sending the + same tag value -- must be unchanged. Two different keys sending the + identical tag value must still share one counter. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 2, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # keyA and keyB share the same bucket -- one call each exhausts the + # shared limit of 2. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}}, + ) + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + + +@pytest.mark.asyncio +async def test_concurrency_scope_by_key_hash_gives_independent_reservations_per_key(time_controller): + """ + scope_by_key_hash=True on a concurrency_limits entry: two different + calling keys sending the identical tag value must not share one + reservation bucket -- keyA exhausting its own single slot must not + block keyB's admission, and releasing keyA's reservation (via the + standard_logging_object.metadata.user_api_key_hash channel) must free + keyA's capacity, not keyB's. Each key is modeled as its own logical + request with its own model_call_details, and keyA's release is spawned + as a genuinely separate child task (mirroring litellm's real dispatch) + to prove release survives that task boundary via the shared + model_call_details object, not via which task happens to run it. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router_scoped_by_key(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + async def _admit(key: str, request_kwargs: dict): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs=request_kwargs, + ) + + async def _release(key: str, kwargs: dict): + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 0, + "response_cost": 0, + "metadata": {"user_api_key_hash": key}, + } + await limiter.async_log_success_event( + kwargs=kwargs, + response_obj=None, + start_time=0, + end_time=0, + ) + + ready_to_release = asyncio.Event() + key_a_request, key_a_kwargs = _call_context(["end_user_id:u1"]) + key_a_request["metadata"]["user_api_key"] = "keyA" + key_b_request, _key_b_kwargs = _call_context(["end_user_id:u1"]) + key_b_request["metadata"]["user_api_key"] = "keyB" + + async def _key_a_admits_then_waits_then_releases_from_the_same_context_chain(): + await _admit("keyA", key_a_request) + await ready_to_release.wait() + await asyncio.create_task(_release("keyA", key_a_kwargs)) + + # keyA occupies its own single slot; keyB, same tag value, different + # key, still admits since it has its own bucket. + key_a_task = asyncio.create_task(_key_a_admits_then_waits_then_releases_from_the_same_context_chain()) + key_b_task = asyncio.create_task(_admit("keyB", key_b_request)) + await key_b_task + # Let key_a_task's admission run up to (but not past) `ready_to_release.wait()`. + await asyncio.sleep(0) + + # keyA is now at its own capacity -- a second keyA request is rejected. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + + # Let keyA's task proceed to its own child-task release. + ready_to_release.set() + await key_a_task + await asyncio.sleep(0) + + # keyA's capacity is freed -- a fresh keyA request now admits. + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyA"}}, + ) + assert result == healthy + + # keyB's own reservation is untouched by keyA's release -- a second keyB + # request is still rejected. If task isolation were broken, keyA's + # release would have drained keyB's reservation too. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}}, + ) + + +# --------------------------------------------------------------------------- +# in-memory cache isolation -- caller-controlled tag buckets must never evict +# the shared cache's other, authentication-bound counters +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_flooding_tag_buckets_does_not_evict_the_shared_cache_authentication_bound_counter( + time_controller, +): + """ + The proxy-wide internal_usage_cache passed into this limiter is also + used by the key/team parallel-request limiter for its own, + authentication-bound counters, and its default InMemoryCache evicts at + 200 items. Without a dedicated in-memory layer for this hook's own + caller-controlled tag buckets, an attacker sending 200+ distinct tag + values could evict an unrelated authentication-bound counter and let + some other caller exceed a limit nothing here configured. + """ + shared_cache = DualCache() + await shared_cache.async_set_cache(key="authentication_bound_counter", value="do-not-evict") + + limiter = _PROXY_ModelBasedTagRateLimitsHook(internal_usage_cache=shared_cache, time_provider=time_controller.now) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1000, "period_seconds": 60}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + for i in range(250): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:flood-{i}"]}}, + ) + + assert await shared_cache.async_get_cache(key="authentication_bound_counter") == "do-not-evict" + + +def _single_request_per_minute_router() -> "litellm.Router": + return litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_max_in_memory_cache_size_setting_lets_high_cardinality_tags_avoid_early_eviction( + time_controller, monkeypatch +): + """ + This hook's own isolated cache still defaults to 200 items, shared across + every distinct tag value it sees. A deployment rate-limiting on a + high-cardinality tag_id (e.g. per end user) without Redis can raise + `litellm_settings.model_based_tag_rate_limits_max_in_memory_cache_size` so an + earlier bucket survives churn from later, unrelated tag values: with + limit=1, a still-live bucket rejects a second request instead of having + been evicted back to a fresh count of 0. + """ + monkeypatch.setattr(litellm, "model_based_tag_rate_limits_max_in_memory_cache_size", 500) + + limiter = _PROXY_ModelBasedTagRateLimitsHook(internal_usage_cache=DualCache(), time_provider=time_controller.now) + router = _single_request_per_minute_router() + limiter.update_variables(llm_router=router) + healthy = router.model_list + + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:early-user"]}}, + ) + + for i in range(250): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:flood-{i}"]}}, + ) + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:early-user"]}}, + ) + + +@pytest.mark.parametrize( + "invalid_configured_size", + [ + 0, # would hit InMemoryCache.set_cache's `max_size_in_memory == 0` short-circuit, disabling the cache + -1, # would loop `heapq.heappop` on an empty heap in InMemoryCache.evict_cache and raise IndexError + "500", # an unresolved os.environ/ substitution or config typo; `len(...) >= "500"` raises TypeError + True, # bool is an int subclass; must not be misread as the positive integer 1 + ], +) +@pytest.mark.asyncio +async def test_invalid_max_in_memory_cache_size_falls_back_to_the_safe_default( + time_controller, monkeypatch, invalid_configured_size +): + """ + DualCache.async_set_cache swallows any exception raised while writing, so an + invalid configured size would otherwise silently disable every counter write + for this hook (every read then sees an empty counter and is admitted) instead + of failing loudly. Each of these must be rejected in favor of the safe + default: a limit=1 bucket must still reject a second, immediate request. + """ + monkeypatch.setattr(litellm, "model_based_tag_rate_limits_max_in_memory_cache_size", invalid_configured_size) + + limiter = _PROXY_ModelBasedTagRateLimitsHook(internal_usage_cache=DualCache(), time_provider=time_controller.now) + router = _single_request_per_minute_router() + limiter.update_variables(llm_router=router) + healthy = router.model_list + + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + +# --------------------------------------------------------------------------- +# per-tag Redis/bucket key TTL override +# --------------------------------------------------------------------------- + + +def _concurrency_limit(period_seconds: int, key_ttl_seconds: int | None = None) -> _ConfiguredLimit: + return _ConfiguredLimit( + unit="concurrency", + entry=TagRateLimitEntry( + name="active", tag_id="end_user_id", limit=1, period_seconds=period_seconds, key_ttl_seconds=key_ttl_seconds + ), + deployment_scope=None, + ) + + +def test_ttl_for_concurrency_honors_key_ttl_seconds_above_the_safety_floor(): + above_floor: Final = _CONCURRENCY_MIN_SAFETY_TTL_SECONDS + 100 + assert ( + _PROXY_ModelBasedTagRateLimitsHook._ttl_for(_concurrency_limit(period_seconds=60, key_ttl_seconds=above_floor)) + == above_floor + ) + + +def test_ttl_for_concurrency_never_drops_below_the_safety_floor_even_with_a_lower_override(): + """ + A reservation's TTL must comfortably outlast any real in-flight request, so + an operator-set override below _CONCURRENCY_MIN_SAFETY_TTL_SECONDS must not + be honored as-is -- a slow request's reservation would otherwise self-heal + (expire) while still genuinely running, silently admitting extra requests. + """ + below_floor: Final = 10 + assert ( + _PROXY_ModelBasedTagRateLimitsHook._ttl_for(_concurrency_limit(period_seconds=5, key_ttl_seconds=below_floor)) + == _CONCURRENCY_MIN_SAFETY_TTL_SECONDS + ) + + +def test_tag_rate_limit_entry_rejects_non_positive_key_ttl_seconds(): + with pytest.raises(ValidationError, match="key_ttl_seconds must be a positive integer"): + TagRateLimitEntry(name="per_minute", limit=1, period_seconds=60, key_ttl_seconds=0) + + +def test_tag_rate_limit_entry_rejects_key_ttl_seconds_shorter_than_period_seconds(): + """ + Regression test for a real bug: a key_ttl_seconds shorter than + period_seconds expires the bucket key before its window rolls over, + resetting the counter to zero mid-window and letting tagged traffic + exceed the configured limit. + """ + with pytest.raises(ValidationError, match="key_ttl_seconds must be at least period_seconds"): + TagRateLimitEntry(name="per_minute", limit=1, period_seconds=60, key_ttl_seconds=59) + + +# --------------------------------------------------------------------------- +# per-tag max_in_memory_cache_size override -- dedicated cache partitions +# --------------------------------------------------------------------------- + + +def test_tag_rate_limit_entry_rejects_non_positive_max_in_memory_cache_size(): + with pytest.raises(ValidationError, match="max_in_memory_cache_size must be a positive integer"): + TagRateLimitEntry(name="per_minute", limit=1, period_seconds=60, max_in_memory_cache_size=0) + + +def _two_request_limit_router(team_limit: int, user_limit: int, user_cache_size: int | None) -> "litellm.Router": + return litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + {"name": "team_cap", "tag_id": "team_id", "limit": team_limit, "period_seconds": 60}, + { + "name": "user_cap", + "tag_id": "end_user_id", + "limit": user_limit, + "period_seconds": 60, + "max_in_memory_cache_size": user_cache_size, + }, + ] + } + }, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_max_in_memory_cache_size_override_isolates_a_flood_on_that_entry_from_a_default_partition_entry( + time_controller, +): + """ + An entry with its own max_in_memory_cache_size gets a dedicated cache + partition. Flooding that entry's own high-cardinality tag values must + never evict a *different* entry's bucket that was never given an + override and still lives on the hook's single default partition. + """ + limiter = _make_limiter(time_controller) + router = _two_request_limit_router(team_limit=1, user_limit=1000, user_cache_size=5) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # team_cap's bucket (default partition) is created and admitted once. + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["team_id:t1"]}} + ) + + # Flood user_cap's own dedicated partition (cap=5) past its own capacity + # many times over -- this must stay fully confined to user_cap's own + # partition and never touch team_cap's default-partition bucket. + for i in range(250): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:flood-{i}"]}}, + ) + + # team_cap's bucket must still be at its limit (1) -- a second team_id:t1 + # request is rejected. If it had been evicted by user_cap's flood, this + # would instead admit (a fresh, zeroed counter). + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["team_id:t1"]}}, + ) + + +@pytest.mark.asyncio +async def test_two_entries_sharing_the_identical_max_in_memory_cache_size_still_get_separate_partitions( + time_controller, +): + """ + Partitions are keyed by the entry's full signature, not the override + value alone: two unrelated entries that happen to choose the identical + max_in_memory_cache_size must not be merged into one shared cache, or + flooding one would evict the other's bucket exactly like the bug this + override exists to fix. + """ + limiter = _make_limiter(time_controller) + # Both team_cap and user_cap set the identical max_in_memory_cache_size (5). + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "team_cap", + "tag_id": "team_id", + "limit": 1, + "period_seconds": 60, + "max_in_memory_cache_size": 5, + }, + { + "name": "user_cap", + "tag_id": "end_user_id", + "limit": 1000, + "period_seconds": 60, + "max_in_memory_cache_size": 5, + }, + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["team_id:t1"]}} + ) + + for i in range(250): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": [f"end_user_id:flood-{i}"]}}, + ) + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["team_id:t1"]}}, + ) + + +@pytest.mark.asyncio +async def test_concurrency_slot_with_a_cache_size_override_is_released_against_the_same_partition(time_controller): + """ + A concurrency reservation on an entry with its own max_in_memory_cache_size + must be released against that same dedicated partition. If the release + path fell back to the default partition instead, it would silently no-op + (nothing to decrement there) and the reservation would leak forever. + """ + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "concurrency_limits": { + "limits": [ + { + "name": "inflight", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 300, + "max_in_memory_cache_size": 10, + } + ] + } + }, + ) + ] + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # At capacity: a second concurrent reservation for the same tag is rejected. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + # The first request completes -- its slot is released against the + # overridden partition -- freeing capacity again. + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 0, + "response_cost": 0, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_token_accounting_with_a_cache_size_override_lands_on_that_entrys_own_partition(time_controller): + """ + tokens/dollars increments go through a per-partition v3 handler (grouped + in async_log_success_event), not always the default one -- an entry with + its own max_in_memory_cache_size must have its usage actually accounted, + not silently dropped or misrouted to the default partition's handler. + """ + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 100, + "period_seconds": 86400, + "max_in_memory_cache_size": 10, + } + ] + } + }, + ) + ] + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 150, + "response_cost": 0, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # 150 tokens already used, over the limit of 100 -- the next admission + # check must reject. If the increment had been silently dropped (never + # reaching the overridden partition), this would incorrectly admit. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + +# --------------------------------------------------------------------------- +# apply_to_key_alias -- shared TagRateLimitEntry field, also usable on a +# per-model entry (the global_tag_rate_limits_hook is its primary motivation, +# but the field composes with async_filter_deployments unmodified) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_restricts_a_per_model_entry_to_the_listed_key(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "per_minute", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # A key with no matching alias is entirely unaffected -- the entry never + # applies to it, so it can call repeatedly with no rejection. + for _ in range(3): + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_alias": "other-key"}}, + ) + assert result == healthy + + # The listed key alias is admitted once, then rejected on its 2nd call. + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_alias": "premium-key"}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_alias": "premium-key"}}, + ) + + +# --------------------------------------------------------------------------- +# apply_to_models -- shared TagRateLimitEntry field, also usable on a +# per-model entry (expected to be rarely useful there, since a per-deployment +# entry is already implicitly scoped to whichever model_name declares it, but +# it must compose identically to every other shared scoping field) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_to_models_ignores_a_non_matching_model_group(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "per_minute", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "apply_to_models": ["other-group"], + } + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + for _ in range(3): + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + +@pytest.mark.asyncio +async def test_apply_to_models_restricts_a_per_model_entry_to_the_listed_model_group(time_controller): + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "request_limits": { + "limits": [ + { + "name": "per_minute", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "apply_to_models": ["grp"], + } + ] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py new file mode 100644 index 00000000000..d645b09075c --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py @@ -0,0 +1,383 @@ +""" +Unit tests for the primitives shared by both tag-scoped rate-limit hooks +(`model_based_tag_rate_limits_hook.py` and `global_tag_rate_limits_hook.py`). + +These test the hook-independent logic in isolation: identity extraction, +`entry_applies` scoping, and the partition/bucket-TTL key helpers. Each +hook's own test file covers everything specific to how it wires these +primitives into its own admission/accounting engine. +""" + +from litellm.proxy.hooks.tag_rate_limits_shared import ( + bucket_ttl_seconds, + entry_applies, + extract_identity, + extract_key_hash, + fixed_length_identity, + order_tags_for_identity_resolution, + partition_key, +) +from litellm.types.router import TagRateLimitEntry, TagRateLimitScope + +# --------------------------------------------------------------------------- +# extract_identity +# --------------------------------------------------------------------------- + + +def test_extract_identity_matches_prefixed_tag(): + assert extract_identity(["team_id:t1", "end_user_id:u1"], "end_user_id") == "u1" + + +def test_extract_identity_returns_none_when_absent(): + assert extract_identity(["team_id:t1"], "end_user_id") is None + + +def test_extract_identity_skips_negation_tags(): + """A `!end_user_id:u1` routing-negation marker must never be read as identity.""" + assert extract_identity(["!end_user_id:u1"], "end_user_id") is None + + +# --------------------------------------------------------------------------- +# order_tags_for_identity_resolution -- veria-ai finding on PR #38292: a +# caller-supplied tag must not shadow a policy-backed (key/team/project) +# tag sharing the same tag_id prefix +# --------------------------------------------------------------------------- + + +def test_order_tags_for_identity_resolution_prefers_inherited_tag_over_caller_supplied(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["company_id:attacker-chosen", "end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "company_id") == "real-company" + + +def test_order_tags_for_identity_resolution_falls_back_to_caller_tags_when_nothing_inherited(): + request_kwargs = {"metadata": {}} + tags = ["end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "end_user_id") == "u1" + + +def test_order_tags_for_identity_resolution_keeps_caller_only_tags_not_shadowed_by_a_different_tag_id(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "end_user_id") == "u1" + assert extract_identity(ordered, "company_id") == "real-company" + + +def test_order_tags_for_identity_resolution_deduplicates_identical_tag_present_in_both_sources(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["company_id:real-company"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert ordered.count("company_id:real-company") == 1 + + +# --------------------------------------------------------------------------- +# fixed_length_identity -- tag_value is caller-controlled with no length +# bound; a hook's own contribution to a cache key must not grow with it +# --------------------------------------------------------------------------- + + +def test_fixed_length_identity_bounds_key_contribution_regardless_of_input_size(): + """ + A caller can submit an arbitrarily long tag value (no length or content + bound is enforced upstream of either hook). Without hashing, that value + would go straight into an in-memory dict key (bypassing + max_in_memory_cache_size, which caps item *count* not key bytes) and an + unbounded-length Redis key (Redis has no key-count or key-size cap at + all here). A fixed-length digest bounds a hook's own contribution to + the key regardless of input size. + """ + huge_value = "x" * 5_000_000 + digest = fixed_length_identity(huge_value) + assert len(digest) == 64 # sha256 hex digest length, independent of input size + + +def test_fixed_length_identity_preserves_distinctness(): + """Hashing must not collapse two different tag values onto one bucket.""" + assert fixed_length_identity("user-a") != fixed_length_identity("user-b") + assert fixed_length_identity("user-a") == fixed_length_identity("user-a") + + +# --------------------------------------------------------------------------- +# extract_key_hash -- must read only the one field the server actually +# authenticates into, never fall back to the other +# --------------------------------------------------------------------------- + + +def test_extract_key_hash_ignores_a_forged_value_in_the_non_authoritative_field(): + """ + On a route where litellm_metadata is authoritative, the server writes + the real hash there and never touches metadata -- so a caller-supplied + metadata.user_api_key must not be read at all, let alone win. + """ + request_kwargs = { + "metadata": {"user_api_key": "forged-by-caller"}, + "litellm_metadata": {"user_api_key": "real-authenticated-hash"}, + } + assert extract_key_hash(request_kwargs, "litellm_metadata") == "real-authenticated-hash" + + +def test_extract_key_hash_reads_metadata_when_it_is_the_authoritative_field(): + request_kwargs = {"metadata": {"user_api_key": "real-hash"}} + assert extract_key_hash(request_kwargs, "metadata") == "real-hash" + + +# --------------------------------------------------------------------------- +# entry_applies -- enabled_for / disabled_for / apply_to_key_alias / apply_to_models +# --------------------------------------------------------------------------- + + +def test_entry_applies_with_none_of_the_scoping_fields_set(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_disabled_for_on_its_own_tag_id_excludes_a_listed_value(): + """disabled_for's `tag_id` can be set to the entry's own tag_id, gating on + a subset of its own resolved identity rather than a second tag.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + assert entry_applies(entry, ["end_user_id:u2"], None, None) is True + + +def test_entry_applies_enabled_for_on_its_own_tag_id_restricts_to_a_listed_value(): + """enabled_for's `tag_id` can likewise be set to the entry's own tag_id, + admitting only a hand-picked subset of its own resolved identity.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u2", "u3")), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + assert entry_applies(entry, ["end_user_id:u2"], None, None) is True + + +def test_entry_applies_matches_an_enabled_for_gate(): + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is True + + +def test_entry_applies_skips_when_enabled_for_gate_tag_is_absent(): + """ + enabled_for is an allowlist gate: absence of the gate tag must not + satisfy it, unlike disabled_for below. + """ + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_skips_when_disabled_for_gate_matches(): + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False + + +def test_entry_applies_when_disabled_for_gate_tag_is_absent(): + """disabled_for is a denylist gate: absence of the gate tag has nothing + to match against, so the entry still applies.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_disabled_for_overrides_a_matching_enabled_for_gate(): + """Deny (disabled_for) takes effect independently of whether the + enabled_for gate itself matched, even when both target the same tag.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False + + +def test_entry_applies_with_apply_to_key_alias_unset_applies_to_every_key(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], "any-key-alias", None) is True + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_admits_a_key_alias_on_the_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], "team-a-key", None) is True + + +def test_entry_applies_rejects_a_key_alias_missing_from_the_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], "team-b-key", None) is False + + +def test_entry_applies_rejects_when_key_has_no_alias_but_allowlist_is_set(): + """apply_to_key_alias is an allowlist gate: a key with no alias at all + never satisfies it, same as enabled_for's absent-gate-tag semantics.""" + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_with_apply_to_models_unset_applies_to_every_model(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_admits_a_model_on_the_apply_to_models_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True + + +def test_entry_applies_rejects_a_model_missing_from_the_apply_to_models_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, "sonnet-chain") is False + + +def test_entry_applies_rejects_when_model_is_absent_but_apply_to_models_is_set(): + """apply_to_models is an allowlist gate: a request with no model at all + never satisfies it, same as apply_to_key_alias's absent-key semantics.""" + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_apply_to_models_composes_with_apply_to_key_alias(): + """Both gates must pass: a request against the listed model but a + non-listed key alias must not apply, even though apply_to_models alone + would have admitted it.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + apply_to_models=("opus-chain",), + apply_to_key_alias=("premium-key",), + ) + assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "opus-chain") is True + assert entry_applies(entry, ["end_user_id:u1"], "other-key", "opus-chain") is False + assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "sonnet-chain") is False + + +# --------------------------------------------------------------------------- +# partition_key -- entries that share max_in_memory_cache_size but disagree +# on any policy-fingerprinted field must never share a cache partition +# --------------------------------------------------------------------------- + + +def test_partition_key_distinguishes_entries_that_differ_only_by_scope_by_key_hash(): + """ + scope_by_key_hash is part of the partition-key signature: two entries + identical in every other field but differing only on this flag are + different rate limits (different bucket keys per each hook's own + `_hash_tag`) and must never be routed to the same cache partition. + """ + unscoped = TagRateLimitEntry( + name="per_minute", tag_id="end_user_id", limit=5, period_seconds=60, max_in_memory_cache_size=100 + ) + scoped = TagRateLimitEntry( + name="per_minute", + tag_id="end_user_id", + limit=5, + period_seconds=60, + scope_by_key_hash=True, + max_in_memory_cache_size=100, + ) + assert partition_key(unscoped) != partition_key(scoped) + + +def test_partition_key_distinguishes_entries_that_differ_only_by_scoping_fields(): + """ + A plain, unscoped entry and a scoped override can legitimately share + name/tag_id/limit/period_seconds/scope_by_key_hash while disagreeing on + enabled_for/disabled_for/apply_to_key_alias/apply_to_models -- + policy_fingerprint already treats that as two distinct policies, so a + shared max_in_memory_cache_size must not route them onto the same + in-memory partition either, or one entry's high-cardinality traffic can + evict the other's active counters from a cache neither entry asked to + share. + """ + base_kwargs = { + "name": "daily", + "tag_id": "end_user_id", + "limit": 100, + "period_seconds": 86400, + "max_in_memory_cache_size": 50, + } + unscoped = TagRateLimitEntry(**base_kwargs) + enabled_for_scoped = TagRateLimitEntry( + **base_kwargs, enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + disabled_for_scoped = TagRateLimitEntry( + **base_kwargs, disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + alias_scoped = TagRateLimitEntry(**base_kwargs, apply_to_key_alias=("premium-key",)) + models_scoped = TagRateLimitEntry(**base_kwargs, apply_to_models=("opus-chain",)) + + keys = { + partition_key(unscoped), + partition_key(enabled_for_scoped), + partition_key(disabled_for_scoped), + partition_key(alias_scoped), + partition_key(models_scoped), + } + assert len(keys) == 5 + + +# --------------------------------------------------------------------------- +# bucket_ttl_seconds -- per-tag Redis/bucket key TTL override +# --------------------------------------------------------------------------- + + +def test_bucket_ttl_seconds_defaults_to_period_plus_one_hour_when_unset(): + entry = TagRateLimitEntry(name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60) + assert bucket_ttl_seconds(entry) == 60 + 3600 + + +def test_bucket_ttl_seconds_honors_key_ttl_seconds_override(): + entry = TagRateLimitEntry( + name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60, key_ttl_seconds=120 + ) + assert bucket_ttl_seconds(entry) == 120 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 64318778bc2..805ee88a136 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -33,6 +33,7 @@ from litellm.proxy.common_request_processing import ( ttft_keepalive_interval, _override_openai_response_model, _parse_event_data_for_error, + _release_disconnect_state_on_all_callbacks, _resolve_per_request_model_group_alias, _should_return_raw_model_name, _UpstreamClosingStreamingResponse, @@ -4063,7 +4064,52 @@ class TestCancelOnDisconnect: llm_call.cancel() with pytest.raises(asyncio.CancelledError): - await _await_llm_call_cancelling_on_disconnect(request, llm_call) + await _await_llm_call_cancelling_on_disconnect(request, llm_call, {}) + + async def test_disconnect_releases_callback_state_before_499(self, monkeypatch): + """ + asyncio.CancelledError is a BaseException, not an Exception, so it + never reaches litellm.utils.wrapper_async's own except block -- the + cancelled call's async_log_failure_event never fires, and the 499 + this raises is later handled by post_call_failure_hook, a different + hook a CustomLogger like model_based_tag_rate_limits_hook doesn't implement. Without + an explicit release here, a callback that reserved per-request state + at admission (a concurrency slot) leaks it until that state's own + safety TTL. This mirrors the streaming disconnect case + (_finalize_streaming_generator_cleanup), just for a non-streaming + call cancelled via the opt-in cancel_on_disconnect flag. + """ + recorder = _RecordingDisconnectHookLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + request = self._request([{"type": "http.disconnect"}]) + llm_call = asyncio.get_running_loop().create_future() + + with pytest.raises(HTTPException) as exc_info: + await _await_llm_call_cancelling_on_disconnect( + request, llm_call, {"litellm_logging_obj": MagicMock()} + ) + + assert exc_info.value.status_code == 499 + assert recorder.disconnect_hook_calls == 1 + + async def test_release_disconnect_state_calls_every_callback_including_ones_without_an_override( + self, monkeypatch + ): + """ + Bugbot finding: CustomLogger never defined async_release_disconnect_state_hook + as an empty default, unlike every other optional hook on that base class, so + calling it on a callback that never overrides it (most registered callbacks) + raised AttributeError -- caught here, but still a real gap in the base class's + own contract that a genuine implementation bug would be indistinguishable from. + """ + overriding = _RecordingDisconnectHookLogger() + bare = CustomLogger() + monkeypatch.setattr(litellm, "callbacks", [overriding, bare]) + + await _release_disconnect_state_on_all_callbacks({"litellm_call_id": "call-1"}) + + assert overriding.disconnect_hook_calls == 1 + assert await bare.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) is None async def _drive_base_process_llm_request( self, monkeypatch, general_settings: dict, llm_call, request: Request @@ -5450,6 +5496,129 @@ 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_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): """ @@ -5596,6 +5765,15 @@ class _RecordingSuccessLogger(CustomLogger): self.success_events.append({"kwargs": kwargs, "response_obj": response_obj}) +class _RecordingDisconnectHookLogger(CustomLogger): + def __init__(self): + super().__init__() + self.disconnect_hook_calls = 0 + + async def async_release_disconnect_state_hook(self, request_data: dict) -> None: + self.disconnect_hook_calls += 1 + + class TestStreamingClientDisconnectBilling: """ A client disconnect throws GeneratorExit into the proxy streaming @@ -5990,6 +6168,64 @@ class TestStreamingClientDisconnectBilling: assert usage.prompt_tokens_details is not None assert usage.prompt_tokens_details.cached_tokens == 7 + @pytest.mark.asyncio + async def test_disconnect_without_billable_chunks_releases_callback_state(self, monkeypatch): + """ + A callback that reserves per-request state outside of the success/failure + logging callbacks (e.g. a concurrency slot admitted before the first + chunk) would otherwise leak it on a disconnect with nothing to bill, + since neither logging callback ever fires for it. The disconnect + cleanup must give every registered callback a chance to release such + state via async_release_disconnect_state_hook. + """ + import types + + response = await self._start_partial_stream() + empty_response = types.SimpleNamespace(chunks=[], messages=None) + recorder = _RecordingDisconnectHookLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=empty_response, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=types.SimpleNamespace( + _arelease_max_parallel_requests_on_disconnect=AsyncMock(), + ), + ) + + assert recorder.disconnect_hook_calls == 1 + + @pytest.mark.asyncio + async def test_disconnect_billing_skips_callback_disconnect_hook(self, monkeypatch): + """ + When a disconnect-time success event already fired (partial billing + dispatched it), that event's own async_log_success_event already ran + for every registered callback. The disconnect hook must not also run + in that case, so a callback with idempotent-but-not-free release logic + does not do redundant work on every disconnect. + """ + import types + + recorder = _RecordingDisconnectHookLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response = await self._start_partial_stream() + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=response, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=types.SimpleNamespace( + _arelease_max_parallel_requests_on_disconnect=AsyncMock(), + ), + ) + + assert recorder.disconnect_hook_calls == 0 + def _apply_stream_usage_tracking( data: dict, diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index accd3b32a0d..f58bd9b118e 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -5,6 +5,7 @@ from litellm.types.router import ( Deployment, LiteLLM_Params, ModelInfo, + TagRateLimitEntry, ) from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams @@ -89,3 +90,47 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") + + +# litellm/types/router.py is imported by plain SDK users, not just the proxy, so a +# TagRateLimitEntry validation error should read as a generic config-validation +# message and not describe the proxy rate-limit hook's internal admission mechanics. +_INTERNAL_ENFORCEMENT_JARGON = ( + "admission", + "tagged request", + "tagged traffic", + "check-and-increment", + "read-only", + "atomic", + "window rolls over", +) + + +def _assert_message_has_no_internal_jargon(excinfo: pytest.ExceptionInfo) -> None: + message = str(excinfo.value).lower() + leaked = [term for term in _INTERNAL_ENFORCEMENT_JARGON if term in message] + assert not leaked, f"validation message leaked internal enforcement jargon: {leaked}" + + +def test_limit_infinite_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=float("inf"), period_seconds=60) + _assert_message_has_no_internal_jargon(excinfo) + + +def test_limit_non_positive_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=0, period_seconds=60) + _assert_message_has_no_internal_jargon(excinfo) + + +def test_key_ttl_seconds_shorter_than_period_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + key_ttl_seconds=60, + ) + _assert_message_has_no_internal_jargon(excinfo) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 959b2eada25..644c0f22063 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16616 + "limit": 16615 }, "LIT011": { "limit": 5583 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 405ec9a01bf..973adce672e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35115,6 +35115,64 @@ export interface components { /** Tpm Limit */ tpm_limit?: number | null; }; + /** TagRateLimitEntry */ + TagRateLimitEntry: { + /** Apply To Key Alias */ + apply_to_key_alias?: string[] | null; + /** Apply To Models */ + apply_to_models?: string[] | null; + disabled_for?: components["schemas"]["TagRateLimitScope"] | null; + enabled_for?: components["schemas"]["TagRateLimitScope"] | null; + /** Key Ttl Seconds */ + key_ttl_seconds?: number | null; + /** Limit */ + limit: number; + /** Max In Memory Cache Size */ + max_in_memory_cache_size?: number | null; + /** Name */ + name: string; + /** Period Seconds */ + period_seconds: number; + /** + * Scope By Key Hash + * @default false + */ + scope_by_key_hash: boolean; + /** + * Tag Id + * @default end_user_id + */ + tag_id: string; + }; + /** TagRateLimitGroup */ + TagRateLimitGroup: { + /** + * Limits + * @default [] + */ + limits: components["schemas"]["TagRateLimitEntry"][]; + }; + /** + * TagRateLimitScope + * @description A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + * entry to `tag_id: company_id, values: ["1032"]` so it only applies to + * requests tagged as belonging to company 1032, independent of whichever + * tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + * `disabled_for`, which are the only two fields that construct this. + */ + TagRateLimitScope: { + /** Tag Id */ + tag_id: string; + /** Values */ + values: string[]; + }; + /** TagRateLimits */ + TagRateLimits: { + concurrency_limits?: components["schemas"]["TagRateLimitGroup"] | null; + dollar_limits?: components["schemas"]["TagRateLimitGroup"] | null; + request_limits?: components["schemas"]["TagRateLimitGroup"] | null; + token_limits?: components["schemas"]["TagRateLimitGroup"] | null; + }; /** * TagSummaryMetrics * @description Summary metrics for a tag @@ -37890,6 +37948,7 @@ export interface components { ptu_effective_from?: string | null; /** Ptu Effective To */ ptu_effective_to?: string | null; + tag_rate_limits?: components["schemas"]["TagRateLimits"] | null; /** Team Id */ team_id?: string | null; /** Team Public Model Name */