From 137b854c507425c3adaa83e84660e7bc411cddb9 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 22:01:46 -0400 Subject: [PATCH] feat(rate-limiting): add per-deployment tag rate limiting hook Enforces token, request, dollar, and concurrency limits scoped to a request tag (end_user_id by default), configured per deployment under model_info.tag_rate_limits and admitted once per routing hop. Supports chain-wide and per-deployment-scoped buckets, team-aliased routing groups, and per-entry scoping via enabled_for/disabled_for/ apply_to_key_alias/apply_to_models. Registers as the model_based_tag_rate_limits_hook callback and reuses the identity extraction, policy fingerprinting, and bucket-key hashing primitives from tag_rate_limits_shared.py. --- litellm/__init__.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 29 + .../hooks/model_based_tag_rate_limits_hook.py | 1642 ++++++ .../test_model_based_tag_rate_limits_hook.py | 4819 +++++++++++++++++ 4 files changed, 6492 insertions(+) create mode 100644 litellm/proxy/hooks/model_based_tag_rate_limits_hook.py create mode 100644 tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py diff --git a/litellm/__init__.py b/litellm/__init__.py index ec2960c196e..f9348f68f1b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -119,6 +119,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "litellm_agent", "dynamic_rate_limiter", "dynamic_rate_limiter_v3", + "model_based_tag_rate_limits_hook", "langsmith", "prometheus", "otel", @@ -391,6 +392,7 @@ 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 model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3018f0c4d24..4f7a7510f96 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4476,6 +4476,26 @@ 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 == "langtrace": if "LANGTRACE_API_KEY" not in os.environ: raise ValueError("LANGTRACE_API_KEY not found in environment variables") @@ -4916,6 +4936,15 @@ 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 == "langtrace": from litellm.integrations.opentelemetry import OpenTelemetry 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..a4532f725f8 --- /dev/null +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -0,0 +1,1642 @@ +""" +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 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 + get_metadata_variable_name_from_kwargs, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # 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 ( + partition_key as _partition_key, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + policy_fingerprint as _policy_fingerprint, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + resolve_success_event_metadata_variable_name as _resolve_success_event_metadata_variable_name, +) +from litellm.proxy.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 team_id this limit was resolved under via `by_team_alias`, or None + # when resolved via `by_model_name`. 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 + # lookup by (team_id, alias). + team_scope: str | None = None + # The real model_name this limit was found under when `resolve()`'s + # direct lookup by the caller-visible model string 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). None when resolved directly, in + # which case the caller-visible name is already unambiguous and safe to + # hash by. Set, this overrides the caller-visible name in the bucket key + # so limits from two different underlying model_names sharing one + # routing group never collide on one counter. + 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 _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 + `get_metadata_variable_name_from_kwargs` 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. + """ + 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(key, replace(limit, resolved_group=name)) # mutable-ok: see docstring above + 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=lambda deployment: deployment["model_name"]) + by_model_name: Final[Mapping[str, tuple[_ConfiguredLimit, ...]]] = MappingProxyType( + { + model_name: configured + for model_name, deployment_group in groupby( + sorted_by_model_name, key=lambda deployment: deployment["model_name"] + ) + 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( + replace(limit, team_scope=alias_key[0]) + 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. +_PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys" + +# 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 +# a caller overwriting their own other request's entry, which only weakens +# that caller's own configured cap rather than crossing between callers. +_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" + + +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 when this + # limit was found via resolve_any()'s per-deployment fallback (routing + # groups): the caller-visible name is ambiguous there (shared by every + # member model_name), so hashing by it would collide two different + # underlying model_names' identically-named limits onto one counter. + # 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: + 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 ContextVar or `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. + """ + 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 + pending.extend(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 + + +@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 + ) -> 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.""" + if self._check_and_incr_script is not None: + raw: Final = await self._check_and_incr_script(keys=(key,), args=(limit, increment, ttl)) + return bool(raw[0]), float(raw[1]) + + async with self._lock: + current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) + current: Final = float(current_value) if current_value is not None else 0.0 + if current + increment > limit: + return False, current + new_value: Final = current + increment + await cache.async_set_cache(key=key, value=new_value, ttl=ttl, litellm_parent_otel_span=None) + return True, new_value + + async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None: + if self._decr_floor_zero_script is not None: + await self._decr_floor_zero_script(keys=(key,), args=(delta,)) + return + async with self._lock: + current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) + current: Final = float(current_value) if current_value is not None else 0.0 + await cache.async_set_cache(key=key, value=max(0.0, current + delta), litellm_parent_otel_span=None) + + async def _atomic_check_and_increment( + self, + checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], + ) -> tuple[int | None, tuple[float, ...]]: + """ + All-or-nothing 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) in enumerate(checks): + admitted = False + try: + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, 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]], up_to_index: int + ) -> None: + for refund_index in range(up_to_index): + refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index] + try: + await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment) + except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest + verbose_proxy_logger.warning( + "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) + metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(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 is + # the only way they're guaranteed to dedup to the identical bucket + # regardless of cooldown state at either point in time. + 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)) + ) + configured: Final = self._index.get(self.llm_router).resolve_any(model, team_id, candidate_model_names) + if not configured: + return healthy_deployments + + tags: Final = _get_tags_from_request_kwargs( + resolved_request_kwargs, metadata_variable_name=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), + ) + 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: + if not isinstance(call_id, str): + return + try: + await self.internal_usage_cache.async_set_cache( + key=_pending_reservations_cache_key(call_id, key_hash), + value=_encode_reservations(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 can only belong to an earlier hop of this same request 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 -- so + admission is never re-entered while an earlier hop's reservation is + still legitimately in flight. + + 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) + 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() + return frozenset(key for key, _partition_key in pending_request_increments) + + async def _pop_pending_concurrency_keys( + self, kwargs: Mapping[str, object] + ) -> tuple[tuple[str, _PartitionKey], ...]: + # Every caller of this method is itself a normal release path, so + # also clear the async_post_call_failure_hook cache mirror for the + # same call_id right here: whatever this pop is about to release + # must never be found there later and double-released. + call_id: Final = kwargs.get("litellm_call_id") + if isinstance(call_id, str): + # 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_success_event_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_success_event_metadata_variable_name(litellm_params_for_metadata) + key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) + try: + await self.internal_usage_cache.dual_cache.async_delete_cache( + _pending_reservations_cache_key(call_id, key_hash) + ) + 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, + ) + # Snapshot then remove only those exact keys, never a blanket clear: + # a sibling hop 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 hop's reservation instead of releasing it later. + pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) + if not isinstance(pending, list) or not pending: + return () + keys: Final = tuple(pending) + for key in keys: + try: + pending.remove(key) + except ValueError: + pass + return keys + + 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. + """ + 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. + 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: + 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_success_event_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). + 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 + ) + 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 ()) + ) + configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names) + if not configured: + return + + tags: Final = _get_tags_from_request_kwargs(kwargs, metadata_variable_name=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=tuple(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/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..897156126c8 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -0,0 +1,4819 @@ +""" +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, + _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 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" + + +# --------------------------------------------------------------------------- +# 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_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"]}, + }, + "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_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"}} + 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"}}, + "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_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 + + +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) + + +# --------------------------------------------------------------------------- +# 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. + """ + 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] + # The alias resolution must carry the team_id into the bucket scope -- + # see test_build_limits_index_keeps_different_teams_same_alias_separate + # for why (two teams can publish the identical alias string). + assert by_name[0].team_scope is None + assert by_alias[0].team_scope == "team-1" + + +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_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) + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: ["key1"]} + + async def detached_release(): + return await limiter._pop_pending_concurrency_keys(model_call_details) + + released = await asyncio.create_task(detached_release()) + assert released == ("key1",) + + # 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"]} + + 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") + return released + + released = await asyncio.create_task(detached_release_then_sibling_admits()) + assert released == ("key1",) + # key2 must still be pending for its own hop's eventual release. + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == ["key2"] + + +@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"]} + first = await limiter._pop_pending_concurrency_keys(model_call_details) + second = await limiter._pop_pending_concurrency_keys(model_call_details) + assert first == ("key1",) + assert second == () + + +# --------------------------------------------------------------------------- +# 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), + (flaky.internal_usage_cache, other_key, 10.0, 1.0, 60), + (flaky.internal_usage_cache, rejecting_key, 0.0, 1.0, 60), + ] + ) + + 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): + if key == raising_key: + raise RuntimeError("simulated transient redis failure") + return await super()._check_and_increment_one(cache, key, limit, increment, 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), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60), + ] + ) + + 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): + 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) + # ...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) + + 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), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60), + ] + ) + + # 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"]}}, + )