From bbaff3695df2c05e8a0760f589c6f5a6a4e5325f Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 21:53:39 -0400 Subject: [PATCH 01/44] feat(types): add TagRateLimitEntry/TagRateLimitScope config types Introduces the tag-scoped rate limit config schema (TagRateLimitEntry, TagRateLimitScope, TagRateLimitGroup, TagRateLimits) and wires it onto ModelInfo.tag_rate_limits, giving tag-based rate limiting hooks a config shape to validate and consume. --- litellm/types/router.py | 177 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) diff --git a/litellm/types/router.py b/litellm/types/router.py index a3335be2b2b..7ab53431755 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum +import math from collections.abc import Mapping from dataclasses import dataclass from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -137,6 +138,180 @@ def _as_utc(value: datetime.datetime | None) -> datetime.datetime | None: return value.astimezone(datetime.timezone.utc) +class TagRateLimitScope(BaseModel): + """ + A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + entry to `tag_id: company_id, values: ["1032"]` so it only applies to + requests tagged as belonging to company 1032, independent of whichever + tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + `disabled_for`, which are the only two fields that construct this. + """ + + tag_id: str + values: tuple[str, ...] + + model_config = ConfigDict(frozen=True) + + @model_validator(mode="after") + def _validate_values(self) -> "TagRateLimitScope": + if not self.values: + raise ValueError("values must be a non-empty list of strings") + return self + + @model_validator(mode="after") + def _normalize_values(self) -> "TagRateLimitScope": + # Sorted and deduplicated: `values` is only ever used for membership + # tests (see _entry_applies), never order-dependent, but is also + # folded verbatim into the dedup signature two deployments' entries + # are compared by (see _scope_signature) -- an unsorted tuple would + # make config-order alone, not policy, decide whether two entries + # dedup to one shared bucket or wrongly split into two. + # object.__setattr__ bypasses this frozen model's own assignment + # guard -- returning a replacement instance from an "after" validator + # is silently ignored when constructing via __init__ (only takes + # effect via model_validate), so mutating in place is the only way + # this normalization reliably applies regardless of construction path. + object.__setattr__(self, "values", tuple(sorted(set(self.values)))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitEntry(BaseModel): + name: str + tag_id: str = "end_user_id" + limit: float + period_seconds: int + scope_by_key_hash: bool = False + # Overrides this entry's bucket/reservation key TTL (Redis, and the + # in-memory fallback when Redis isn't configured). Defaults to + # period_seconds + 3600 when unset -- see _PROXY_ModelBasedTagRateLimitsHook._ttl_for. + # A high-cardinality tag_id can keep many keys alive at once; lowering + # this lets an operator shed them sooner without shortening + # period_seconds itself. + key_ttl_seconds: int | None = None + # Overrides the size of the dedicated in-memory cache partition this + # entry's own keys live in, when Redis isn't configured (or as a local + # fast-path cache when it is). Unset means this entry shares the hook's + # single default partition, sized by + # litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200 if that's also + # unset). A high-cardinality tag_id can churn past that shared cap and + # evict another entry's active counters; setting this gives the entry + # its own dedicated partition instead. + max_in_memory_cache_size: int | None = None + # Gate this entry on a tag -- often a SECOND, independent tag (e.g. + # `enabled_for: {tag_id: company_id, values: ["1032"]}` to scope an + # override to one company's traffic), but `tag_id` can equally be set to + # this same entry's own `tag_id` to scope by a subset of its own + # resolved identity instead, without a second tag at all. + # `disabled_for` is checked first (deny overrides allow) when both are + # set. An absent gate tag never satisfies `enabled_for` (an allowlist + # gate requires an explicit match) but never triggers `disabled_for` + # either (nothing to match against a denylist). + enabled_for: TagRateLimitScope | None = None + disabled_for: TagRateLimitScope | None = None + # Restrict this entry to requests authenticated with one of these virtual + # keys' own `key_alias`. Unset (the default) means the entry applies to + # every request regardless of which key made it. A key with no alias set + # never satisfies this allowlist, same "absent gate never matches an + # allowlist" precedent as `enabled_for`. + apply_to_key_alias: tuple[str, ...] | None = None + # Restrict this entry to requests whose caller-facing `model` matches one + # of these names. Unset (the default) means the entry applies to every + # model. A request with no `model` field never satisfies this allowlist, + # same "absent gate never matches an allowlist" precedent as + # `apply_to_key_alias`. + apply_to_models: tuple[str, ...] | None = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _validate_limit(self) -> "TagRateLimitEntry": + # NaN compares False against every ordering operator, so a NaN limit + # makes the atomic requests/concurrency check-and-increment (which + # rejects when the new value exceeds the limit) never reject -- + # admitting indefinitely -- while the read-only tokens/dollars check + # (which admits when the current value is under the limit) never + # admits, rejecting every tagged request. Either outcome silently + # defeats the entry; reject it at config load time instead. + if math.isnan(self.limit): + raise ValueError("limit must not be NaN") + if math.isinf(self.limit): + raise ValueError( + "limit must be finite -- positive infinity makes admission never reject (current + increment " + "> limit is always false), negative infinity makes it always reject every tagged request" + ) + if self.limit <= 0: + raise ValueError( + "limit must be a positive number -- zero 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, silently blocking all matching traffic instead of the " + "likely intended config" + ) + return self + + @model_validator(mode="after") + def _validate_period_seconds(self) -> "TagRateLimitEntry": + if self.period_seconds <= 0: + raise ValueError("period_seconds must be a positive integer") + return self + + @model_validator(mode="after") + def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": + if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: + raise ValueError("key_ttl_seconds must be a positive integer when set") + if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: + raise ValueError( + "key_ttl_seconds must be at least period_seconds when set -- a shorter TTL expires the " + "counter before its window rolls over, letting tagged traffic reset to zero and exceed the limit" + ) + return self + + @model_validator(mode="after") + def _validate_max_in_memory_cache_size(self) -> "TagRateLimitEntry": + if self.max_in_memory_cache_size is not None and self.max_in_memory_cache_size <= 0: + raise ValueError("max_in_memory_cache_size must be a positive integer when set") + return self + + @model_validator(mode="after") + def _validate_apply_to_key_alias(self) -> "TagRateLimitEntry": + if self.apply_to_key_alias is not None and not self.apply_to_key_alias: + raise ValueError("apply_to_key_alias must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_key_alias(self) -> "TagRateLimitEntry": + # Sorted and deduplicated for the same reason as + # TagRateLimitScope._normalize_values: only ever used for membership + # tests, but also folded verbatim into the dedup signature, where an + # unsorted tuple would make config-order alone decide whether two + # deployments' entries dedup to one shared bucket. + if self.apply_to_key_alias is not None: + self.apply_to_key_alias = tuple(sorted(set(self.apply_to_key_alias))) # mutable-ok: frozen before escaping + return self + + @model_validator(mode="after") + def _validate_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None and not self.apply_to_models: + raise ValueError("apply_to_models must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None: + self.apply_to_models = tuple(sorted(set(self.apply_to_models))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitGroup(BaseModel): + limits: tuple[TagRateLimitEntry, ...] = () + + +class TagRateLimits(BaseModel): + token_limits: TagRateLimitGroup | None = None + request_limits: TagRateLimitGroup | None = None + dollar_limits: TagRateLimitGroup | None = None + concurrency_limits: TagRateLimitGroup | None = None + + class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. @@ -186,6 +361,8 @@ class ModelInfo(MirroredPricingParams): # router-wide default. enable_tag_filtering: bool | None = None + tag_rate_limits: TagRateLimits | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided From 22ad42e37d3474d785a431428d74a09ab2d5071c Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 21:53:52 -0400 Subject: [PATCH 02/44] feat(rate-limiting): extract shared tag-rate-limit helpers into their own module Both tag-scoped rate limiting hooks need the same identity/scope extraction, policy fingerprinting, bucket-key hashing, and cache partitioning primitives. Moving them into their own module lets a model-independent global hook consume them without reaching into a model-based hook's private internals, which is how the two hooks previously shared this logic. --- litellm/proxy/hooks/tag_rate_limits_shared.py | 334 +++++++++++++++++ .../hooks/test_tag_rate_limits_shared.py | 346 ++++++++++++++++++ 2 files changed, 680 insertions(+) create mode 100644 litellm/proxy/hooks/tag_rate_limits_shared.py create mode 100644 tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py new file mode 100644 index 00000000000..5f48f3cf1c1 --- /dev/null +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -0,0 +1,334 @@ +""" +Primitives shared by both tag-scoped rate-limit hooks: +`model_based_tag_rate_limits_hook.py` (per-deployment limits nested under +`model_info.tag_rate_limits`, admitted once per routing hop) and +`global_tag_rate_limits_hook.py` (model-independent limits declared once +under `litellm_settings.global_tag_rate_limits`, admitted once per request). + +Both hooks enforce the identical `TagRateLimitEntry` shape and need the same +identity/scope extraction, policy fingerprinting, bucket-key hashing, and +cache-partitioning primitives, so those live here rather than in either +hook's own module -- keeping neither hook reaching into the other's private +internals to reuse them. Every name here is this module's own public +interface (no leading underscore): each hook imports what it needs aliased +back to its own historical, underscore-prefixed local name (e.g. +`entry_applies as _entry_applies`), so this is a genuine export, not a +private symbol either hook reaches across a module boundary to grab. +""" + +import asyncio +import hashlib +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from litellm.exceptions import RateLimitType +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import TagRateLimitEntry, TagRateLimitScope + +LimitUnit: TypeAlias = Literal["tokens", "requests", "dollars", "concurrency"] +LIMIT_UNITS: Final[tuple[LimitUnit, ...]] = ("tokens", "requests", "dollars", "concurrency") + +# Units whose admission must be atomic (check-and-increment in one Redis +# round trip) because the increment amount is known upfront (always 1). +# tokens/dollars can't be: real usage is only known after the response, so +# they stay a read-then-account-on-success check with a documented, +# unavoidable admit-vs-account race. +ATOMIC_UNITS: Final[frozenset[LimitUnit]] = frozenset({"requests", "concurrency"}) + +UNIT_TO_GROUP_FIELD: Final[Mapping[LimitUnit, str]] = MappingProxyType( + { + "tokens": "token_limits", + "requests": "request_limits", + "dollars": "dollar_limits", + "concurrency": "concurrency_limits", + } +) +UNIT_TO_RATE_LIMIT_TYPE: Final[Mapping[LimitUnit, RateLimitType]] = MappingProxyType( + { + "tokens": RateLimitType.TOKENS, + "requests": RateLimitType.REQUESTS, + "dollars": RateLimitType.BUDGET, + "concurrency": RateLimitType.CONCURRENT_REQUESTS, + } +) + +# Shared read-only fallback for an absent/None mapping (request_kwargs, +# metadata, model_info, ...): avoids constructing a fresh mutable `{}` at +# every one of these call sites just to immediately call `.get()` on it. +EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) + +# `asyncio.create_task`'s own docs: "Save a reference to the result of this +# function, to avoid a task disappearing mid-execution. The event loop only +# keeps weak references to tasks. A task that isn't referenced elsewhere may +# get garbage collected at any time, even before it's done." The success path +# deliberately fires-and-forgets its concurrency release and its token/dollar +# accounting increment (unlike the failure/disconnect paths, which await +# concurrency release directly) to keep the hot success-response path from +# waiting on a Redis round trip; by the time either background task would +# run, the state it needs (popped pending keys, or the request's own usage +# figures) is only available in that task's own closure, so a collected +# task's work is unrecoverable, not just delayed. Holding a strong reference +# here until each task's own completion callback discards it is the standard +# fix, shared by every fire-and-forget task either hook creates. +BACKGROUND_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: see comment above + +# Floor for a concurrency reservation's self-heal TTL, regardless of the +# configured period_seconds. A reservation that expires while its request is +# still genuinely in flight silently admits requests past the limit; this +# generous floor keeps that window far larger than any realistic request +# duration, at the cost of a leaked (crashed-worker) slot self-healing more +# slowly. period_seconds can still raise the TTL further, never lower it. +CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600 + +# Single-key atomic check-and-increment. Deliberately one key per script call +# (never a batch of differently-hash-tagged keys in one call): every tag_rl +# key carries its own self-contained {..} hash tag so unrelated buckets never +# forcibly co-locate on the same Redis Cluster shard, which means a single Lua +# invocation can never span more than one key's slot without risking a +# cross-slot error. All-or-nothing across a hop's multiple atomic checks +# (e.g. requests + concurrency checked together) is achieved in Python by +# calling this once per key and refunding every earlier admission in the same +# batch if a later one is rejected -- the same refund-on-rollback shape as +# `atomic_check_and_increment_by_n` in parallel_request_limiter_v3.py, applied +# per-key instead of per-descriptor since each key already is one hash-tag +# group by construction. +TAG_RL_CHECK_AND_INCR_SCRIPT: Final = """ +local key = KEYS[1] +local limit = tonumber(ARGV[1]) +local increment = tonumber(ARGV[2]) +local ttl = tonumber(ARGV[3]) +local current = tonumber(redis.call('GET', key) or 0) +if current + increment > limit then + return { 0, current } +end +local new_value = redis.call('INCRBY', key, increment) +local current_ttl = redis.call('TTL', key) +if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', key, ttl) +end +return { 1, new_value } +""" + +# Atomic decrement that never leaves a counter negative. Used both to refund +# an earlier admission when a later key in the same batch is rejected, and to +# release a concurrency reservation -- floors at 0 so a decrement that can't +# be attributed to the exact reservation that caused it (see +# `_release_keys`'s docstring) degrades to under-counting rather than a +# negative counter that would admit unlimited requests. Floors via DEL, not +# `SET key 0`: releasing a reservation whose key already expired makes +# INCRBY recreate it with no TTL, and a plain SET would leave that recreated +# key permanently in Redis (SET clears any TTL); DEL removes it outright, +# which reads back identically to 0 everywhere this key is read (`GET key or 0`). +TAG_RL_DECR_FLOOR_ZERO_SCRIPT: Final = """ +local key = KEYS[1] +local delta = tonumber(ARGV[1]) +local new_value = redis.call('INCRBY', key, delta) +if new_value < 0 then + redis.call('DEL', key) + new_value = 0 +end +return new_value +""" + +# A (tag_id, values) pair mirroring TagRateLimitScope's own fields, used to +# fold `enabled_for`/`disabled_for` into a hashable form (a policy +# fingerprint, or a per-deployment dedup signature) without depending on +# TagRateLimitScope's own hashability. +ScopeSignature: TypeAlias = tuple[str, tuple[str, ...]] | None + + +def scope_signature(scope: TagRateLimitScope | None) -> ScopeSignature: + """Normalizes a `TagRateLimitScope` into a plain, hashable tuple, so + `enabled_for`/`disabled_for` can be folded into a policy fingerprint (or + a per-deployment dedup signature) -- two entries disagreeing on either + field must be treated as genuinely different policies rather than + merged into one bucket.""" + return None if scope is None else (scope.tag_id, scope.values) + + +def extract_identity(tags: Sequence[str], tag_id: str) -> str | None: + """ + First tag matching `f"{tag_id}:"`, value after the colon. Tags starting + with `!` are tag-routing negation markers, not identity tags, and are + skipped so they can never be misread as an identity value. + """ + prefix: Final = f"{tag_id}:" + for tag in tags: + if tag.startswith("!"): + continue + if tag.startswith(prefix): + return tag[len(prefix) :] + return None + + +def entry_applies(entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, model: str | None) -> bool: + """ + Applies `entry`'s own scoping fields (`enabled_for`/`disabled_for`/ + `apply_to_key_alias`/`apply_to_models`), evaluated in this order -- deny + overrides allow, checked before any allowlist: + + 1. `disabled_for`: the gate tag (often a SECOND, independent tag, but + `disabled_for.tag_id` can equally be set to this entry's own + `tag_id` to gate on a subset of its own resolved identity) is + present and its value is in `disabled_for.values` -> doesn't apply. + Absent gate tag never triggers this -- nothing to match against a + denylist. + 2. `enabled_for`: the gate tag is absent, or present but its value is + NOT in `enabled_for.values` -> doesn't apply. Unlike `disabled_for`, + absence DOES fail this check -- an allowlist gate requires an + explicit match, so "not tagged at all" means "not in scope". + 3. `apply_to_models`: `model` is absent, or present but not in the + list -> doesn't apply. Same allowlist semantics as `enabled_for` -- + a request with no `model` never satisfies this gate. + 4. `apply_to_key_alias`: the calling key's own alias is absent, or + present but not in the list -> doesn't apply. Same allowlist + semantics as `enabled_for` -- a key with no alias set never + satisfies this gate. + + An entry with none of these fields set always applies -- this is the + unscoped behavior every existing entry has today, unchanged. + """ + if entry.disabled_for is not None: + disabled_gate_value: Final = extract_identity(tags, entry.disabled_for.tag_id) + if disabled_gate_value is not None and disabled_gate_value in entry.disabled_for.values: + return False + if entry.enabled_for is not None: + enabled_gate_value: Final = extract_identity(tags, entry.enabled_for.tag_id) + if enabled_gate_value is None or enabled_gate_value not in entry.enabled_for.values: + return False + if entry.apply_to_models is not None and model not in entry.apply_to_models: + return False + if entry.apply_to_key_alias is None: + return True + return key_alias in entry.apply_to_key_alias + + +def resolve_success_event_metadata_variable_name( + litellm_params_for_metadata: Mapping[str, object], +) -> Literal["metadata", "litellm_metadata"]: + """`get_metadata_variable_name_from_kwargs` only checks key presence, which + misresolves at `async_log_success_event` time: `kwargs["litellm_params"]` + always carries a `litellm_metadata` key (typically `None`) alongside the + real, populated `metadata` dict for a standard (non + LITELLM_METADATA_ROUTES) request, so the key-presence check always picks + `litellm_metadata` there and silently reads no tags/identity at all. + Requiring the value to actually be a populated dict, matching + `_get_request_tags`'s own truthiness check in litellm_logging.py, only + ever prefers `litellm_metadata` when it is genuinely the field the proxy + wrote identity/tags into (LITELLM_METADATA_ROUTES pre-seed it before + admission runs, so it is always a populated dict by success time there).""" + litellm_metadata: Final = litellm_params_for_metadata.get("litellm_metadata") + if isinstance(litellm_metadata, Mapping) and litellm_metadata: + return "litellm_metadata" + return "metadata" + + +def extract_key_hash(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Same single-authoritative-field lookup as + `model_based_tag_rate_limits_hook._extract_team_id`, but for the calling + virtual key's hash: `LiteLLMProxyRequestSetup` sets + `metadata["user_api_key"]` to `user_api_key_dict.api_key`, which despite + the plain name is already the hashed token (see `litellm_pre_call_utils.py`). + """ + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + key_hash: Final = active.get("user_api_key") + return key_hash if isinstance(key_hash, str) else None + + +def extract_key_alias(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Same single-authoritative-field lookup as + `model_based_tag_rate_limits_hook._extract_team_id`, but for the calling + virtual key's own `key_alias`: `LiteLLMProxyRequestSetup` sets + `metadata["user_api_key_alias"]` to `user_api_key_dict.key_alias` + (see `litellm_pre_call_utils.py`).""" + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + key_alias: Final = active.get("user_api_key_alias") + return key_alias if isinstance(key_alias, str) else None + + +def fixed_length_identity(tag_value: str) -> str: + """ + `tag_value` is caller-controlled (whatever follows the tag_id prefix in + a caller-supplied tag) with no length or content bound. Embedding it + directly would let a caller inflate a hook's own in-memory dict keys + past what `max_in_memory_cache_size` bounds (that caps item *count*, not + key bytes) and grow unbounded Redis keys with no cap at all. Hashing to + a fixed-length digest bounds a hook's own contribution to key size + regardless of the caller's input, while still preserving distinctness + (two different tag values still resolve to two different buckets). + """ + return hashlib.sha256(tag_value.encode()).hexdigest() + + +def policy_fingerprint(entry: TagRateLimitEntry) -> str: + """ + Two entries can share a `name` and `tag_id` while genuinely disagreeing + on `limit`, `period_seconds`, `scope_by_key_hash`, or any of the scoping + fields -- both hooks already treat that as two distinct policies + elsewhere (model_based_tag_rate_limits_hook.py's own per-deployment + dedup signature is one example), so the Redis/in-memory bucket key must + too, or two differently-configured entries that happen to share a name + check and charge the identical counter. `scope_by_key_hash` specifically + needs its own slot here rather than relying on a hook's own `_hash_tag` + `key_hash`-derived suffix to carry it: that suffix is empty whenever + `key_hash` resolves to `None` (no virtual key on the call), which would + otherwise collide an unscoped entry with a key-hash-scoped one that + agrees on every other field. Hashed to a fixed-length digest for the + same reason `fixed_length_identity` hashes `tag_value`: an operator's + own `enabled_for`/`disabled_for`/`apply_to_key_alias`/`apply_to_models` + list has no length bound. + """ + fingerprint_source: Final = ( + entry.limit, + entry.period_seconds, + entry.scope_by_key_hash, + scope_signature(entry.enabled_for), + scope_signature(entry.disabled_for), + entry.apply_to_key_alias, + entry.apply_to_models, + ) + return hashlib.sha256(repr(fingerprint_source).encode()).hexdigest()[:16] + + +def bucket_ttl_seconds(entry: TagRateLimitEntry) -> int: + """Redis (and in-memory fallback) TTL for a non-concurrency bucket key. + `entry.key_ttl_seconds` overrides the default of period_seconds + 3600 + when set -- see TagRateLimitEntry.key_ttl_seconds.""" + return entry.key_ttl_seconds if entry.key_ttl_seconds is not None else entry.period_seconds + 3600 + + +# None => this entry shares the hook's single default cache partition +# (matching every entry's behavior before this override existed). Otherwise +# a value-stable signature -- not the override int alone -- so two different +# entries that happen to choose the identical max_in_memory_cache_size don't +# get merged into one shared partition; the same entry (same config content) +# always resolves to the same signature across config/index rebuilds, which +# is what keeps each hook's own `_partitions` cache from leaking a fresh +# partition every time its configuration is re-resolved. +# The str before the trailing int is `policy_fingerprint(entry)`: two +# entries can share tag_id/name while genuinely disagreeing on +# limit/period_seconds/scope_by_key_hash/enabled_for/disabled_for/ +# apply_to_key_alias/apply_to_models -- both hooks already treat that as two +# distinct policies for bucket-key purposes, so a shared +# max_in_memory_cache_size must not route them onto the same partition +# either, or one entry's high-cardinality traffic can evict the other's +# active counters from a cache neither entry asked to share. +# max_in_memory_cache_size stays the trailing element: `partition_key[-1]` +# reads it directly to size the partition's cache. +PartitionKey: TypeAlias = tuple[str, str, str, int] | None +# Grouping type for async_log_success_event's per-partition tokens/dollars +# pipeline dispatch -- named only so the declaration fits on one line; see +# that method for why the grouping is needed. +PartitionOperations: TypeAlias = dict[PartitionKey, list[RedisPipelineIncrementOperation]] + + +def partition_key(entry: TagRateLimitEntry) -> PartitionKey: + if entry.max_in_memory_cache_size is None: + return None + return ( + entry.tag_id, + entry.name, + policy_fingerprint(entry), + entry.max_in_memory_cache_size, + ) diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py new file mode 100644 index 00000000000..153737ecda1 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py @@ -0,0 +1,346 @@ +""" +Unit tests for the primitives shared by both tag-scoped rate-limit hooks +(`model_based_tag_rate_limits_hook.py` and `global_tag_rate_limits_hook.py`). + +These test the hook-independent logic in isolation: identity extraction, +`entry_applies` scoping, and the partition/bucket-TTL key helpers. Each +hook's own test file covers everything specific to how it wires these +primitives into its own admission/accounting engine. +""" + +from litellm.proxy.hooks.tag_rate_limits_shared import ( + bucket_ttl_seconds, + entry_applies, + extract_identity, + extract_key_hash, + fixed_length_identity, + partition_key, +) +from litellm.types.router import TagRateLimitEntry, TagRateLimitScope + +# --------------------------------------------------------------------------- +# extract_identity +# --------------------------------------------------------------------------- + + +def test_extract_identity_matches_prefixed_tag(): + assert extract_identity(["team_id:t1", "end_user_id:u1"], "end_user_id") == "u1" + + +def test_extract_identity_returns_none_when_absent(): + assert extract_identity(["team_id:t1"], "end_user_id") is None + + +def test_extract_identity_skips_negation_tags(): + """A `!end_user_id:u1` routing-negation marker must never be read as identity.""" + assert extract_identity(["!end_user_id:u1"], "end_user_id") is None + + +# --------------------------------------------------------------------------- +# fixed_length_identity -- tag_value is caller-controlled with no length +# bound; a hook's own contribution to a cache key must not grow with it +# --------------------------------------------------------------------------- + + +def test_fixed_length_identity_bounds_key_contribution_regardless_of_input_size(): + """ + A caller can submit an arbitrarily long tag value (no length or content + bound is enforced upstream of either hook). Without hashing, that value + would go straight into an in-memory dict key (bypassing + max_in_memory_cache_size, which caps item *count* not key bytes) and an + unbounded-length Redis key (Redis has no key-count or key-size cap at + all here). A fixed-length digest bounds a hook's own contribution to + the key regardless of input size. + """ + huge_value = "x" * 5_000_000 + digest = fixed_length_identity(huge_value) + assert len(digest) == 64 # sha256 hex digest length, independent of input size + + +def test_fixed_length_identity_preserves_distinctness(): + """Hashing must not collapse two different tag values onto one bucket.""" + assert fixed_length_identity("user-a") != fixed_length_identity("user-b") + assert fixed_length_identity("user-a") == fixed_length_identity("user-a") + + +# --------------------------------------------------------------------------- +# extract_key_hash -- must read only the one field the server actually +# authenticates into, never fall back to the other +# --------------------------------------------------------------------------- + + +def test_extract_key_hash_ignores_a_forged_value_in_the_non_authoritative_field(): + """ + On a route where litellm_metadata is authoritative, the server writes + the real hash there and never touches metadata -- so a caller-supplied + metadata.user_api_key must not be read at all, let alone win. + """ + request_kwargs = { + "metadata": {"user_api_key": "forged-by-caller"}, + "litellm_metadata": {"user_api_key": "real-authenticated-hash"}, + } + assert extract_key_hash(request_kwargs, "litellm_metadata") == "real-authenticated-hash" + + +def test_extract_key_hash_reads_metadata_when_it_is_the_authoritative_field(): + request_kwargs = {"metadata": {"user_api_key": "real-hash"}} + assert extract_key_hash(request_kwargs, "metadata") == "real-hash" + + +# --------------------------------------------------------------------------- +# entry_applies -- enabled_for / disabled_for / apply_to_key_alias / apply_to_models +# --------------------------------------------------------------------------- + + +def test_entry_applies_with_none_of_the_scoping_fields_set(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_disabled_for_on_its_own_tag_id_excludes_a_listed_value(): + """disabled_for's `tag_id` can be set to the entry's own tag_id, gating on + a subset of its own resolved identity rather than a second tag.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + assert entry_applies(entry, ["end_user_id:u2"], None, None) is True + + +def test_entry_applies_enabled_for_on_its_own_tag_id_restricts_to_a_listed_value(): + """enabled_for's `tag_id` can likewise be set to the entry's own tag_id, + admitting only a hand-picked subset of its own resolved identity.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u2", "u3")), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + assert entry_applies(entry, ["end_user_id:u2"], None, None) is True + + +def test_entry_applies_matches_an_enabled_for_gate(): + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is True + + +def test_entry_applies_skips_when_enabled_for_gate_tag_is_absent(): + """ + enabled_for is an allowlist gate: absence of the gate tag must not + satisfy it, unlike disabled_for below. + """ + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_skips_when_disabled_for_gate_matches(): + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False + + +def test_entry_applies_when_disabled_for_gate_tag_is_absent(): + """disabled_for is a denylist gate: absence of the gate tag has nothing + to match against, so the entry still applies.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_disabled_for_overrides_a_matching_enabled_for_gate(): + """Deny (disabled_for) takes effect independently of whether the + enabled_for gate itself matched, even when both target the same tag.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False + + +def test_entry_applies_with_apply_to_key_alias_unset_applies_to_every_key(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], "any-key-alias", None) is True + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_admits_a_key_alias_on_the_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], "team-a-key", None) is True + + +def test_entry_applies_rejects_a_key_alias_missing_from_the_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], "team-b-key", None) is False + + +def test_entry_applies_rejects_when_key_has_no_alias_but_allowlist_is_set(): + """apply_to_key_alias is an allowlist gate: a key with no alias at all + never satisfies it, same as enabled_for's absent-gate-tag semantics.""" + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_with_apply_to_models_unset_applies_to_every_model(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_admits_a_model_on_the_apply_to_models_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True + + +def test_entry_applies_rejects_a_model_missing_from_the_apply_to_models_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, "sonnet-chain") is False + + +def test_entry_applies_rejects_when_model_is_absent_but_apply_to_models_is_set(): + """apply_to_models is an allowlist gate: a request with no model at all + never satisfies it, same as apply_to_key_alias's absent-key semantics.""" + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_apply_to_models_composes_with_apply_to_key_alias(): + """Both gates must pass: a request against the listed model but a + non-listed key alias must not apply, even though apply_to_models alone + would have admitted it.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + apply_to_models=("opus-chain",), + apply_to_key_alias=("premium-key",), + ) + assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "opus-chain") is True + assert entry_applies(entry, ["end_user_id:u1"], "other-key", "opus-chain") is False + assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "sonnet-chain") is False + + +# --------------------------------------------------------------------------- +# partition_key -- entries that share max_in_memory_cache_size but disagree +# on any policy-fingerprinted field must never share a cache partition +# --------------------------------------------------------------------------- + + +def test_partition_key_distinguishes_entries_that_differ_only_by_scope_by_key_hash(): + """ + scope_by_key_hash is part of the partition-key signature: two entries + identical in every other field but differing only on this flag are + different rate limits (different bucket keys per each hook's own + `_hash_tag`) and must never be routed to the same cache partition. + """ + unscoped = TagRateLimitEntry( + name="per_minute", tag_id="end_user_id", limit=5, period_seconds=60, max_in_memory_cache_size=100 + ) + scoped = TagRateLimitEntry( + name="per_minute", + tag_id="end_user_id", + limit=5, + period_seconds=60, + scope_by_key_hash=True, + max_in_memory_cache_size=100, + ) + assert partition_key(unscoped) != partition_key(scoped) + + +def test_partition_key_distinguishes_entries_that_differ_only_by_scoping_fields(): + """ + A plain, unscoped entry and a scoped override can legitimately share + name/tag_id/limit/period_seconds/scope_by_key_hash while disagreeing on + enabled_for/disabled_for/apply_to_key_alias/apply_to_models -- + policy_fingerprint already treats that as two distinct policies, so a + shared max_in_memory_cache_size must not route them onto the same + in-memory partition either, or one entry's high-cardinality traffic can + evict the other's active counters from a cache neither entry asked to + share. + """ + base_kwargs = { + "name": "daily", + "tag_id": "end_user_id", + "limit": 100, + "period_seconds": 86400, + "max_in_memory_cache_size": 50, + } + unscoped = TagRateLimitEntry(**base_kwargs) + enabled_for_scoped = TagRateLimitEntry( + **base_kwargs, enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + disabled_for_scoped = TagRateLimitEntry( + **base_kwargs, disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + alias_scoped = TagRateLimitEntry(**base_kwargs, apply_to_key_alias=("premium-key",)) + models_scoped = TagRateLimitEntry(**base_kwargs, apply_to_models=("opus-chain",)) + + keys = { + partition_key(unscoped), + partition_key(enabled_for_scoped), + partition_key(disabled_for_scoped), + partition_key(alias_scoped), + partition_key(models_scoped), + } + assert len(keys) == 5 + + +# --------------------------------------------------------------------------- +# bucket_ttl_seconds -- per-tag Redis/bucket key TTL override +# --------------------------------------------------------------------------- + + +def test_bucket_ttl_seconds_defaults_to_period_plus_one_hour_when_unset(): + entry = TagRateLimitEntry(name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60) + assert bucket_ttl_seconds(entry) == 60 + 3600 + + +def test_bucket_ttl_seconds_honors_key_ttl_seconds_override(): + entry = TagRateLimitEntry( + name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60, key_ttl_seconds=120 + ) + assert bucket_ttl_seconds(entry) == 120 From f613b13d02797b72ae7fd26a400fbccc5425afbb Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 22:09:05 -0400 Subject: [PATCH 03/44] chore(ui): regenerate schema.d.ts for TagRateLimits types Keeps the dashboard's generated API types in sync with the new TagRateLimitEntry/TagRateLimitScope/TagRateLimitGroup/TagRateLimits schema on ModelInfo. --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 25fbd53018a..9a968e53fde 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35110,6 +35110,64 @@ export interface components { /** Tpm Limit */ tpm_limit?: number | null; }; + /** TagRateLimitEntry */ + TagRateLimitEntry: { + /** Apply To Key Alias */ + apply_to_key_alias?: string[] | null; + /** Apply To Models */ + apply_to_models?: string[] | null; + disabled_for?: components["schemas"]["TagRateLimitScope"] | null; + enabled_for?: components["schemas"]["TagRateLimitScope"] | null; + /** Key Ttl Seconds */ + key_ttl_seconds?: number | null; + /** Limit */ + limit: number; + /** Max In Memory Cache Size */ + max_in_memory_cache_size?: number | null; + /** Name */ + name: string; + /** Period Seconds */ + period_seconds: number; + /** + * Scope By Key Hash + * @default false + */ + scope_by_key_hash: boolean; + /** + * Tag Id + * @default end_user_id + */ + tag_id: string; + }; + /** TagRateLimitGroup */ + TagRateLimitGroup: { + /** + * Limits + * @default [] + */ + limits: components["schemas"]["TagRateLimitEntry"][]; + }; + /** + * TagRateLimitScope + * @description A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + * entry to `tag_id: company_id, values: ["1032"]` so it only applies to + * requests tagged as belonging to company 1032, independent of whichever + * tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + * `disabled_for`, which are the only two fields that construct this. + */ + TagRateLimitScope: { + /** Tag Id */ + tag_id: string; + /** Values */ + values: string[]; + }; + /** TagRateLimits */ + TagRateLimits: { + concurrency_limits?: components["schemas"]["TagRateLimitGroup"] | null; + dollar_limits?: components["schemas"]["TagRateLimitGroup"] | null; + request_limits?: components["schemas"]["TagRateLimitGroup"] | null; + token_limits?: components["schemas"]["TagRateLimitGroup"] | null; + }; /** * TagSummaryMetrics * @description Summary metrics for a tag @@ -37885,6 +37943,7 @@ export interface components { ptu_effective_from?: string | null; /** Ptu Effective To */ ptu_effective_to?: string | null; + tag_rate_limits?: components["schemas"]["TagRateLimits"] | null; /** Team Id */ team_id?: string | null; /** Team Public Model Name */ From 86739e2366be9a008bdfb9ff5a0db02470626721 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:01:32 -0400 Subject: [PATCH 04/44] fix(types): keep TagRateLimitEntry validation messages proxy-agnostic litellm/types/router.py is imported by plain SDK users, not just the proxy; the previous ValueError messages for limit/key_ttl_seconds explained the proxy rate-limit hook's internal admission mechanics (atomic check-and-increment, read-only tokens/dollars check, cache TTL rollover), leaking implementation details across the SDK/proxy boundary. Move that mechanistic reasoning into code comments for future maintainers and keep the raised messages generic, per Greptile's finding on PR #38289. Adds regression tests asserting the three affected validators reject their invalid inputs without leaking proxy-internal enforcement jargon. --- litellm/types/router.py | 24 ++++++------- tests/test_litellm/types/test_router.py | 45 +++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 14 deletions(-) diff --git a/litellm/types/router.py b/litellm/types/router.py index 7ab53431755..47eae2bf4a8 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -234,18 +234,15 @@ class TagRateLimitEntry(BaseModel): # defeats the entry; reject it at config load time instead. if math.isnan(self.limit): raise ValueError("limit must not be NaN") + # Positive infinity never rejects the checks that gate this limit; negative + # infinity always does. Both silently defeat the entry. if math.isinf(self.limit): - raise ValueError( - "limit must be finite -- positive infinity makes admission never reject (current + increment " - "> limit is always false), negative infinity makes it always reject every tagged request" - ) + raise ValueError("limit must be finite") + # Zero or negative makes every check that gates this limit either always + # reject or never admit, silently blocking or admitting all matching traffic + # instead of the likely intended config. if self.limit <= 0: - raise ValueError( - "limit must be a positive number -- zero 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, silently blocking all matching traffic instead of the " - "likely intended config" - ) + raise ValueError("limit must be a positive number") return self @model_validator(mode="after") @@ -258,11 +255,10 @@ class TagRateLimitEntry(BaseModel): def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: raise ValueError("key_ttl_seconds must be a positive integer when set") + # A shorter TTL than period_seconds expires the counter before its period + # elapses, letting it reset early and exceed the limit. if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: - raise ValueError( - "key_ttl_seconds must be at least period_seconds when set -- a shorter TTL expires the " - "counter before its window rolls over, letting tagged traffic reset to zero and exceed the limit" - ) + raise ValueError("key_ttl_seconds must be at least period_seconds when set") return self @model_validator(mode="after") diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index accd3b32a0d..f58bd9b118e 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -5,6 +5,7 @@ from litellm.types.router import ( Deployment, LiteLLM_Params, ModelInfo, + TagRateLimitEntry, ) from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams @@ -89,3 +90,47 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") + + +# litellm/types/router.py is imported by plain SDK users, not just the proxy, so a +# TagRateLimitEntry validation error should read as a generic config-validation +# message and not describe the proxy rate-limit hook's internal admission mechanics. +_INTERNAL_ENFORCEMENT_JARGON = ( + "admission", + "tagged request", + "tagged traffic", + "check-and-increment", + "read-only", + "atomic", + "window rolls over", +) + + +def _assert_message_has_no_internal_jargon(excinfo: pytest.ExceptionInfo) -> None: + message = str(excinfo.value).lower() + leaked = [term for term in _INTERNAL_ENFORCEMENT_JARGON if term in message] + assert not leaked, f"validation message leaked internal enforcement jargon: {leaked}" + + +def test_limit_infinite_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=float("inf"), period_seconds=60) + _assert_message_has_no_internal_jargon(excinfo) + + +def test_limit_non_positive_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=0, period_seconds=60) + _assert_message_has_no_internal_jargon(excinfo) + + +def test_key_ttl_seconds_shorter_than_period_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + key_ttl_seconds=60, + ) + _assert_message_has_no_internal_jargon(excinfo) From 6862980ea31e11be930962ac0edd52fc3cbf09c1 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:28:38 -0400 Subject: [PATCH 05/44] chore: retrigger CI (previous run's jobs were cancelled by infra/concurrency, not a real failure) From 529104d65fd2ef2ef8e5a01b784b1eecb15b661f Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:36:53 -0400 Subject: [PATCH 06/44] fix(proxy): stop a caller-supplied tag from shadowing a policy-backed identity tag extract_identity/entry_applies (this module's own functions) resolve a tag_id via first-match-by-prefix over metadata.tags, but _merge_tags (litellm_pre_call_utils.py) keeps caller-supplied tags ahead of key/team/ project tags in that merged list. An authenticated caller could submit e.g. company_id:attacker-chosen ahead of the calling key's real company_id:real-company tag and have every rate-limit entry scoped to company_id resolve to the caller's own value instead of the key's. Adds order_tags_for_identity_resolution, which puts metadata.inherited_tags (the server-computed snapshot of only the tags the calling key/team/project's own config contributed) ahead of the full tags list before either lookup runs. veria-ai caught this while reviewing #38292 (whose branch currently carries this module's commits); porting the fix here since the vulnerable functions it defends are this PR's own. #38292 will wire the call sites in once it rebases onto this branch instead of carrying its own duplicate copy. --- litellm/proxy/hooks/tag_rate_limits_shared.py | 22 +++++++++++ .../hooks/test_tag_rate_limits_shared.py | 37 +++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index 5f48f3cf1c1..34989e5bf5a 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -247,6 +247,28 @@ def extract_key_alias(request_kwargs: Mapping[str, object], metadata_variable_na return key_alias if isinstance(key_alias, str) else None +def order_tags_for_identity_resolution( + tags: Sequence[str], request_kwargs: Mapping[str, object], metadata_variable_name: str +) -> tuple[str, ...]: + """`extract_identity`/`entry_applies` both resolve a `tag_id` via + first-match-by-prefix. `_merge_tags` (litellm_pre_call_utils.py) appends + key/team/project tags only if not already present, keeping caller-supplied + tags first in the merged `tags` list -- so an authenticated caller could + submit e.g. `company_id:attacker-chosen` ahead of the calling key's real + `company_id:real-company` tag and have every entry scoped to `company_id` + resolve to the caller's own value instead of the key's. `metadata.inherited_tags` + is a separate, server-computed snapshot of only the tags the calling + key/team/project's own config contributed (see that field's docstring in + litellm_pre_call_utils.py), so putting it first makes a policy-backed tag + win over a same-prefix caller-supplied one. + """ + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + inherited_tags: Final = active.get("inherited_tags") if isinstance(active, Mapping) else None + if not isinstance(inherited_tags, (list, tuple)) or not inherited_tags: + return tuple(tags) + return tuple(dict.fromkeys((*inherited_tags, *tags))) + + def fixed_length_identity(tag_value: str) -> str: """ `tag_value` is caller-controlled (whatever follows the tag_id prefix in diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py index 153737ecda1..d645b09075c 100644 --- a/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py @@ -14,6 +14,7 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( extract_identity, extract_key_hash, fixed_length_identity, + order_tags_for_identity_resolution, partition_key, ) from litellm.types.router import TagRateLimitEntry, TagRateLimitScope @@ -36,6 +37,42 @@ def test_extract_identity_skips_negation_tags(): assert extract_identity(["!end_user_id:u1"], "end_user_id") is None +# --------------------------------------------------------------------------- +# order_tags_for_identity_resolution -- veria-ai finding on PR #38292: a +# caller-supplied tag must not shadow a policy-backed (key/team/project) +# tag sharing the same tag_id prefix +# --------------------------------------------------------------------------- + + +def test_order_tags_for_identity_resolution_prefers_inherited_tag_over_caller_supplied(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["company_id:attacker-chosen", "end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "company_id") == "real-company" + + +def test_order_tags_for_identity_resolution_falls_back_to_caller_tags_when_nothing_inherited(): + request_kwargs = {"metadata": {}} + tags = ["end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "end_user_id") == "u1" + + +def test_order_tags_for_identity_resolution_keeps_caller_only_tags_not_shadowed_by_a_different_tag_id(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "end_user_id") == "u1" + assert extract_identity(ordered, "company_id") == "real-company" + + +def test_order_tags_for_identity_resolution_deduplicates_identical_tag_present_in_both_sources(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["company_id:real-company"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert ordered.count("company_id:real-company") == 1 + + # --------------------------------------------------------------------------- # fixed_length_identity -- tag_value is caller-controlled with no length # bound; a hook's own contribution to a cache key must not grow with it From f8de40d3d08190f76a41c2da24540dcd7351b4e6 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:57:13 -0400 Subject: [PATCH 07/44] chore: retrigger CI (zizmor cancelled by infra/concurrency at the queue stage, not a real failure) From ac9f2abccad4d60c24de56541379a8fca6ad7c08 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 21:53:39 -0400 Subject: [PATCH 08/44] feat(types): add TagRateLimitEntry/TagRateLimitScope config types Introduces the tag-scoped rate limit config schema (TagRateLimitEntry, TagRateLimitScope, TagRateLimitGroup, TagRateLimits) and wires it onto ModelInfo.tag_rate_limits, giving tag-based rate limiting hooks a config shape to validate and consume. --- litellm/types/router.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/litellm/types/router.py b/litellm/types/router.py index 47eae2bf4a8..7ab53431755 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -234,15 +234,18 @@ class TagRateLimitEntry(BaseModel): # defeats the entry; reject it at config load time instead. if math.isnan(self.limit): raise ValueError("limit must not be NaN") - # Positive infinity never rejects the checks that gate this limit; negative - # infinity always does. Both silently defeat the entry. if math.isinf(self.limit): - raise ValueError("limit must be finite") - # Zero or negative makes every check that gates this limit either always - # reject or never admit, silently blocking or admitting all matching traffic - # instead of the likely intended config. + raise ValueError( + "limit must be finite -- positive infinity makes admission never reject (current + increment " + "> limit is always false), negative infinity makes it always reject every tagged request" + ) if self.limit <= 0: - raise ValueError("limit must be a positive number") + raise ValueError( + "limit must be a positive number -- zero 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, silently blocking all matching traffic instead of the " + "likely intended config" + ) return self @model_validator(mode="after") @@ -255,10 +258,11 @@ class TagRateLimitEntry(BaseModel): def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: raise ValueError("key_ttl_seconds must be a positive integer when set") - # A shorter TTL than period_seconds expires the counter before its period - # elapses, letting it reset early and exceed the limit. if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: - raise ValueError("key_ttl_seconds must be at least period_seconds when set") + raise ValueError( + "key_ttl_seconds must be at least period_seconds when set -- a shorter TTL expires the " + "counter before its window rolls over, letting tagged traffic reset to zero and exceed the limit" + ) return self @model_validator(mode="after") From 10d7aff0b66738e330e7ffa4746ec336ff80527b Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:01:32 -0400 Subject: [PATCH 09/44] fix(types): keep TagRateLimitEntry validation messages proxy-agnostic litellm/types/router.py is imported by plain SDK users, not just the proxy; the previous ValueError messages for limit/key_ttl_seconds explained the proxy rate-limit hook's internal admission mechanics (atomic check-and-increment, read-only tokens/dollars check, cache TTL rollover), leaking implementation details across the SDK/proxy boundary. Move that mechanistic reasoning into code comments for future maintainers and keep the raised messages generic, per Greptile's finding on PR #38289. Adds regression tests asserting the three affected validators reject their invalid inputs without leaking proxy-internal enforcement jargon. --- litellm/types/router.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/litellm/types/router.py b/litellm/types/router.py index 7ab53431755..47eae2bf4a8 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -234,18 +234,15 @@ class TagRateLimitEntry(BaseModel): # defeats the entry; reject it at config load time instead. if math.isnan(self.limit): raise ValueError("limit must not be NaN") + # Positive infinity never rejects the checks that gate this limit; negative + # infinity always does. Both silently defeat the entry. if math.isinf(self.limit): - raise ValueError( - "limit must be finite -- positive infinity makes admission never reject (current + increment " - "> limit is always false), negative infinity makes it always reject every tagged request" - ) + raise ValueError("limit must be finite") + # Zero or negative makes every check that gates this limit either always + # reject or never admit, silently blocking or admitting all matching traffic + # instead of the likely intended config. if self.limit <= 0: - raise ValueError( - "limit must be a positive number -- zero 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, silently blocking all matching traffic instead of the " - "likely intended config" - ) + raise ValueError("limit must be a positive number") return self @model_validator(mode="after") @@ -258,11 +255,10 @@ class TagRateLimitEntry(BaseModel): def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: raise ValueError("key_ttl_seconds must be a positive integer when set") + # A shorter TTL than period_seconds expires the counter before its period + # elapses, letting it reset early and exceed the limit. if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: - raise ValueError( - "key_ttl_seconds must be at least period_seconds when set -- a shorter TTL expires the " - "counter before its window rolls over, letting tagged traffic reset to zero and exceed the limit" - ) + raise ValueError("key_ttl_seconds must be at least period_seconds when set") return self @model_validator(mode="after") From 137b854c507425c3adaa83e84660e7bc411cddb9 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 22:01:46 -0400 Subject: [PATCH 10/44] 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"]}}, + ) From 2b99f4af3379173f6001a1b7858b8ef7dfb7edd2 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 22:01:57 -0400 Subject: [PATCH 11/44] fix(proxy): release rate limit hook state on client disconnect A client disconnect throws GeneratorExit/CancelledError into the request path, so neither the success nor failure logging callback runs and a concurrency slot reserved at admission leaks until its own safety TTL. Gives every registered CustomLogger a chance to release such state via the new async_release_disconnect_state_hook, called from both the streaming and non-streaming cancel-on-disconnect paths. --- litellm/proxy/common_request_processing.py | 54 +++++++++-- .../proxy/test_common_request_processing.py | 95 ++++++++++++++++++- 2 files changed, 138 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 315fbcba310..09d52b78da8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,10 +29,10 @@ from litellm.constants import ( NON_INFERENCE_CALL_TYPES, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, - STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -203,10 +203,6 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } -def _withheld_provider_output(response: object) -> bool: - return getattr(response, "has_buffered_provider_output", False) is True - - def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -397,6 +393,32 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return True +async def _release_disconnect_state_on_all_callbacks(request_data: Mapping[str, object]) -> None: + """ + A client disconnect throws GeneratorExit/CancelledError into the streaming + generator, so neither the success nor failure logging callback runs for it + (see the callers of this function). A callback that reserves per-request + state outside of those two callbacks (e.g. a concurrency slot admitted + before the first chunk) would otherwise leak that state until its own + safety-net TTL. Give every registered callback a chance to release such + state via the optional, default-no-op ``async_release_disconnect_state_hook``. + + Only ``CustomLogger`` instances are considered, never raw string entries: + by the time a request can reach this proxy-only cleanup path, startup's + ``ProxyLogging._init_litellm_callbacks`` has already replaced every string + entry in ``litellm.callbacks`` with its initialized instance in place. + """ + for callback in litellm.callbacks: + if not isinstance(callback, CustomLogger): + continue + try: + await callback.async_release_disconnect_state_hook(request_data) + except Exception as e: # noqa: BLE001 # one callback's cleanup must never block another's or the response teardown + verbose_proxy_logger.debug( + "Failed to run async_release_disconnect_state_hook for %s: %s", type(callback).__name__, e + ) + + async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None: pending_tasks: Final = [task for task in tasks if not task.done()] for task in pending_tasks: @@ -1454,6 +1476,7 @@ async def _cancel_llm_call_on_client_disconnect( async def _await_llm_call_cancelling_on_disconnect( request: Request, llm_api_call: "asyncio.Future[_LlmCallT]", + request_data: Mapping[str, object], ) -> _LlmCallT: disconnect_event: Final = asyncio.Event() monitor: Final = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event)) @@ -1461,6 +1484,14 @@ async def _await_llm_call_cancelling_on_disconnect( return await llm_api_call except asyncio.CancelledError: if disconnect_event.is_set(): + # This cancellation never reaches litellm.utils.wrapper_async's own + # except block (asyncio.CancelledError is a BaseException, not an + # Exception, since Python 3.8), so async_log_failure_event never + # fires for it -- the same gap async_release_disconnect_state_hook + # was added for on the streaming path (see + # _finalize_streaming_generator_cleanup), just reached here via a + # cancelled non-streaming call instead of a mid-stream disconnect. + await _release_disconnect_state_on_all_callbacks(request_data) raise HTTPException( status_code=499, detail=_CLIENT_DISCONNECT_DETAIL, @@ -2285,7 +2316,9 @@ class ProxyBaseLLMRequestProcessing: try: if general_settings.get("cancel_on_disconnect", False): - responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) + responses = await _await_llm_call_cancelling_on_disconnect( # rebind-ok: assigned in exactly one of these two mutually exclusive branches + request, llm_responses, self.data + ) else: responses = await llm_responses finally: @@ -3364,6 +3397,8 @@ class ProxyBaseLLMRequestProcessing: and user_api_key_dict is not None ): await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + if not success_event_owns_slot_release: + await _release_disconnect_state_on_all_callbacks(request_data) if hasattr(response, "aclose"): try: @@ -3447,9 +3482,8 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. A keepalive ping carries no provider output, - # so it must not suppress that refund. - delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES + # False and refunds. + delivered_chunk = True yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): @@ -3463,7 +3497,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk and not _withheld_provider_output(response): + if not delivered_chunk: from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 64318778bc2..f854d8f94e5 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4063,7 +4063,33 @@ class TestCancelOnDisconnect: llm_call.cancel() with pytest.raises(asyncio.CancelledError): - await _await_llm_call_cancelling_on_disconnect(request, llm_call) + await _await_llm_call_cancelling_on_disconnect(request, llm_call, {}) + + async def test_disconnect_releases_callback_state_before_499(self, monkeypatch): + """ + asyncio.CancelledError is a BaseException, not an Exception, so it + never reaches litellm.utils.wrapper_async's own except block -- the + cancelled call's async_log_failure_event never fires, and the 499 + this raises is later handled by post_call_failure_hook, a different + hook a CustomLogger like model_based_tag_rate_limits_hook doesn't implement. Without + an explicit release here, a callback that reserved per-request state + at admission (a concurrency slot) leaks it until that state's own + safety TTL. This mirrors the streaming disconnect case + (_finalize_streaming_generator_cleanup), just for a non-streaming + call cancelled via the opt-in cancel_on_disconnect flag. + """ + recorder = _RecordingDisconnectHookLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + request = self._request([{"type": "http.disconnect"}]) + llm_call = asyncio.get_running_loop().create_future() + + with pytest.raises(HTTPException) as exc_info: + await _await_llm_call_cancelling_on_disconnect( + request, llm_call, {"litellm_logging_obj": MagicMock()} + ) + + assert exc_info.value.status_code == 499 + assert recorder.disconnect_hook_calls == 1 async def _drive_base_process_llm_request( self, monkeypatch, general_settings: dict, llm_call, request: Request @@ -5596,6 +5622,15 @@ class _RecordingSuccessLogger(CustomLogger): self.success_events.append({"kwargs": kwargs, "response_obj": response_obj}) +class _RecordingDisconnectHookLogger(CustomLogger): + def __init__(self): + super().__init__() + self.disconnect_hook_calls = 0 + + async def async_release_disconnect_state_hook(self, request_data: dict) -> None: + self.disconnect_hook_calls += 1 + + class TestStreamingClientDisconnectBilling: """ A client disconnect throws GeneratorExit into the proxy streaming @@ -5990,6 +6025,64 @@ class TestStreamingClientDisconnectBilling: assert usage.prompt_tokens_details is not None assert usage.prompt_tokens_details.cached_tokens == 7 + @pytest.mark.asyncio + async def test_disconnect_without_billable_chunks_releases_callback_state(self, monkeypatch): + """ + A callback that reserves per-request state outside of the success/failure + logging callbacks (e.g. a concurrency slot admitted before the first + chunk) would otherwise leak it on a disconnect with nothing to bill, + since neither logging callback ever fires for it. The disconnect + cleanup must give every registered callback a chance to release such + state via async_release_disconnect_state_hook. + """ + import types + + response = await self._start_partial_stream() + empty_response = types.SimpleNamespace(chunks=[], messages=None) + recorder = _RecordingDisconnectHookLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=empty_response, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=types.SimpleNamespace( + _arelease_max_parallel_requests_on_disconnect=AsyncMock(), + ), + ) + + assert recorder.disconnect_hook_calls == 1 + + @pytest.mark.asyncio + async def test_disconnect_billing_skips_callback_disconnect_hook(self, monkeypatch): + """ + When a disconnect-time success event already fired (partial billing + dispatched it), that event's own async_log_success_event already ran + for every registered callback. The disconnect hook must not also run + in that case, so a callback with idempotent-but-not-free release logic + does not do redundant work on every disconnect. + """ + import types + + recorder = _RecordingDisconnectHookLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + response = await self._start_partial_stream() + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=response, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=types.SimpleNamespace( + _arelease_max_parallel_requests_on_disconnect=AsyncMock(), + ), + ) + + assert recorder.disconnect_hook_calls == 0 + def _apply_stream_usage_tracking( data: dict, From b88fa3bf367aacac979e50ba0bec5e263b86d3c7 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:15:21 -0400 Subject: [PATCH 12/44] fix(proxy): restore keepalive-ping exclusion from streaming disconnect refund The disconnect-state-release hook added in the previous commit dropped the existing has_buffered_provider_output guard and the STREAM_SSE_KEEPALIVE_PING_BYTES exclusion while rewiring the streaming generator's cleanup path, so a client disconnecting after only keepalive pings (or while an agentic stream holds back real output) got refunded to input cost even when billable output had already been generated. Restores both checks; veria-ai caught this on review, and the existing test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_cost regression test now passes again. --- litellm/proxy/common_request_processing.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 09d52b78da8..b582b164609 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,6 +29,7 @@ from litellm.constants import ( NON_INFERENCE_CALL_TYPES, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, + STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -203,6 +204,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } +def _withheld_provider_output(response: object) -> bool: + return getattr(response, "has_buffered_provider_output", False) is True + + def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -3482,8 +3487,9 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. - delivered_chunk = True + # False and refunds. A keepalive ping carries no provider output, + # so it must not suppress that refund. + delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): @@ -3497,7 +3503,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk: + if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, ) From d1cc806ba40c4b98066d26d5ff26806b13835e46 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:15:34 -0400 Subject: [PATCH 13/44] fix(proxy): stop a caller-supplied tag from shadowing a policy-backed identity tag extract_identity/entry_applies resolve a tag_id via first-match-by-prefix over metadata.tags, but _merge_tags keeps caller-supplied tags ahead of key/team/ project tags in that merged list. An authenticated caller could submit e.g. company_id:attacker-chosen ahead of the calling key's real company_id:real-company tag and have every rate-limit entry scoped to company_id resolve to the caller's own value instead of the key's. Adds order_tags_for_identity_resolution, which puts metadata.inherited_tags (the server-computed snapshot of only the tags the calling key/team/project's own config contributed) ahead of the full tags list before either lookup runs, and wires it into both call sites in model_based_tag_rate_limits_hook.py. veria-ai caught this on review. --- .../hooks/model_based_tag_rate_limits_hook.py | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index a4532f725f8..2f1cf21f3af 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -89,6 +89,9 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( from litellm.proxy.hooks.tag_rate_limits_shared import ( fixed_length_identity as _fixed_length_identity, ) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + order_tags_for_identity_resolution as _order_tags_for_identity_resolution, +) from litellm.proxy.hooks.tag_rate_limits_shared import ( partition_key as _partition_key, ) @@ -573,8 +576,16 @@ _PENDING_REQUEST_INCREMENTS_FIELD: Final[str] = "_model_based_tag_rate_limits_pe # 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. +# one caller reusing its own call_id across two of its own concurrent +# requests. This is not confined to "weakens only that caller's own cap": +# for a chain-wide (non scope_by_key_hash) tag, the released reservation is +# on a bucket that tag value's other callers share too, so the forging +# caller's own terminal failure can release a slot on a bucket a different +# caller is also drawing from. Closing this fully needs a per-admission +# identifier that is both server-generated (unlike litellm_call_id) and +# survives proxy/utils.py's post_call_failure_hook stripping +# litellm_logging_obj (unlike everything this comment already ruled out +# above) -- tracked as a known follow-up rather than attempted here. _PENDING_RESERVATIONS_CACHE_KEY_PREFIX: Final = "model_based_tag_rate_limits:pending_reservations:" @@ -1080,8 +1091,10 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if not configured: return healthy_deployments - tags: Final = _get_tags_from_request_kwargs( - resolved_request_kwargs, metadata_variable_name=metadata_variable_name + tags: Final = _order_tags_for_identity_resolution( + _get_tags_from_request_kwargs(resolved_request_kwargs, metadata_variable_name=metadata_variable_name), + resolved_request_kwargs, + metadata_variable_name, ) present_deployment_ids: Final[frozenset[str]] = frozenset( @@ -1593,7 +1606,11 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if not configured: return - tags: Final = _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name) + tags: Final = _order_tags_for_identity_resolution( + _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name), + kwargs, + metadata_variable_name, + ) if not tags: return From 3a2f893134a310ed4fb3013b51d7bf14d0401d41 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 14:01:32 -0400 Subject: [PATCH 14/44] fix(proxy): resolve inherited_tags from litellm_params at success-event time too order_tags_for_identity_resolution only checked the top level of the metadata dict, which is correct for admission's flat request_kwargs but never present at async_log_success_event time -- Logging.model_call_details only ever nests metadata under kwargs["litellm_params"]. Admission correctly preferred the key-backed identity tag, but token/dollar accounting fell through to the caller-forged one instead, charging a different bucket than the one admission actually checked. Adds the same litellm_params fallback _get_tags_from_request_kwargs already relies on. bugbot caught this on review. --- litellm/proxy/hooks/tag_rate_limits_shared.py | 24 +++++++- .../test_model_based_tag_rate_limits_hook.py | 57 +++++++++++++++++++ 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index 34989e5bf5a..ad8e7bfe29d 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -247,6 +247,26 @@ def extract_key_alias(request_kwargs: Mapping[str, object], metadata_variable_na return key_alias if isinstance(key_alias, str) else None +def _active_metadata_bucket(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> Mapping[str, object]: + """Same fallback `_get_tags_from_request_kwargs` (tag_based_routing.py) + already relies on: `request_kwargs` is a flat, top-level-metadata dict at + admission time, but `Logging.model_call_details` (what `kwargs` actually + is by `async_log_success_event`/`async_log_failure_event` time) never + carries `metadata`/`litellm_metadata` at its own top level, only nested + under `request_kwargs["litellm_params"]`. Checking only the top level + silently finds nothing at success/failure time, exactly the same failure + mode that lookup already had to handle.""" + top_level: Final = request_kwargs.get(metadata_variable_name) + if isinstance(top_level, Mapping): + return top_level + litellm_params: Final = request_kwargs.get("litellm_params") + if isinstance(litellm_params, Mapping): + nested: Final = litellm_params.get(metadata_variable_name) + if isinstance(nested, Mapping): + return nested + return EMPTY_MAPPING + + def order_tags_for_identity_resolution( tags: Sequence[str], request_kwargs: Mapping[str, object], metadata_variable_name: str ) -> tuple[str, ...]: @@ -262,8 +282,8 @@ def order_tags_for_identity_resolution( litellm_pre_call_utils.py), so putting it first makes a policy-backed tag win over a same-prefix caller-supplied one. """ - active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING - inherited_tags: Final = active.get("inherited_tags") if isinstance(active, Mapping) else None + active: Final = _active_metadata_bucket(request_kwargs, metadata_variable_name) + inherited_tags: Final = active.get("inherited_tags") if not isinstance(inherited_tags, (list, tuple)) or not inherited_tags: return tuple(tags) return tuple(dict.fromkeys((*inherited_tags, *tags))) 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 index 897156126c8..33edf7519a3 100644 --- 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 @@ -1378,6 +1378,63 @@ async def test_log_success_event_accounts_when_litellm_params_carries_a_null_lit ) +@pytest.mark.asyncio +async def test_log_success_event_accounts_the_key_backed_tag_not_a_caller_forged_one(time_controller): + """ + Bugbot finding: admission (async_filter_deployments) sees a flat + request_kwargs where metadata.inherited_tags sits at the top level, but + kwargs at async_log_success_event time is Logging.model_call_details, + which only ever nests metadata under kwargs["litellm_params"] (see + test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key). + order_tags_for_identity_resolution's own inherited_tags lookup only + checked the top level, so a caller-forged company_id tag that admission + correctly ignored could still get accounted against at success time, + charging a different bucket than the one admission actually checked. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "company_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "metadata": { + "tags": ["company_id:attacker-chosen"], + "inherited_tags": ["company_id:real-company"], + }, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + real_key = _expected_bucket_key("grp", "tokens", "daily", "company_id", "real-company", 86400, now, limit=500000) + forged_key = _expected_bucket_key( + "grp", "tokens", "daily", "company_id", "attacker-chosen", 86400, now, limit=500000 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=real_key, litellm_parent_otel_span=None)) == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=forged_key, litellm_parent_otel_span=None) is None + + @pytest.mark.asyncio async def test_log_success_event_reads_nested_litellm_metadata_when_that_is_authoritative(time_controller): """ From 54a7f6f7b319ccfb10a2fea480175c76362831da Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 14:01:43 -0400 Subject: [PATCH 15/44] fix(logging): add async_release_disconnect_state_hook as a default no-op on CustomLogger Every other optional hook on CustomLogger ships as an empty method a subclass can override; this one didn't, so _release_disconnect_state_on_all_callbacks calling it on any callback that doesn't implement it (nearly all of them) raised AttributeError, caught and debug-logged on every single disconnect. bugbot caught this on review. --- litellm/integrations/custom_logger.py | 10 ++++++++++ .../proxy/test_common_request_processing.py | 20 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 41caf732db0..139fc2b2e65 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -187,6 +187,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): pass + async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + """ + Called when a client disconnects mid-request, for a callback that reserved + per-request state outside async_log_success_event/async_log_failure_event + (e.g. a concurrency slot admitted before the first response chunk) -- those + two callbacks never run for a client disconnect, so a callback relying on + them alone to release such state would otherwise leak it until its own + safety-net TTL. + """ + async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"): """Called when an audit log is created. Override in subclasses to handle.""" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f854d8f94e5..e092a6d6272 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -33,6 +33,7 @@ from litellm.proxy.common_request_processing import ( ttft_keepalive_interval, _override_openai_response_model, _parse_event_data_for_error, + _release_disconnect_state_on_all_callbacks, _resolve_per_request_model_group_alias, _should_return_raw_model_name, _UpstreamClosingStreamingResponse, @@ -4091,6 +4092,25 @@ class TestCancelOnDisconnect: assert exc_info.value.status_code == 499 assert recorder.disconnect_hook_calls == 1 + async def test_release_disconnect_state_calls_every_callback_including_ones_without_an_override( + self, monkeypatch + ): + """ + Bugbot finding: CustomLogger never defined async_release_disconnect_state_hook + as an empty default, unlike every other optional hook on that base class, so + calling it on a callback that never overrides it (most registered callbacks) + raised AttributeError -- caught here, but still a real gap in the base class's + own contract that a genuine implementation bug would be indistinguishable from. + """ + overriding = _RecordingDisconnectHookLogger() + bare = CustomLogger() + monkeypatch.setattr(litellm, "callbacks", [overriding, bare]) + + await _release_disconnect_state_on_all_callbacks({"litellm_call_id": "call-1"}) + + assert overriding.disconnect_hook_calls == 1 + assert await bare.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) is None + async def _drive_base_process_llm_request( self, monkeypatch, general_settings: dict, llm_call, request: Request ): From 513c2737501b971e699ebaca73213030167e2d54 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 15:08:56 -0400 Subject: [PATCH 16/44] fix(proxy): refresh a concurrency key's Redis TTL on every admission, not just its first TAG_RL_CHECK_AND_INCR_SCRIPT only called EXPIRE when a key had no TTL at all, so a concurrency counter's expiry was fixed from its first admission and never pushed out by later ones. A concurrency bucket isn't epoch-windowed like requests/tokens/dollars -- its TTL exists purely as a crash-safety net for a reservation whose explicit release never runs -- so a still-active bucket under sustained traffic would expire mid-flight, silently admitting past the cap and letting a later release decrement an unrelated, newer cohort's counter. Adds a refresh_ttl script argument, true only for the concurrency caller, and verified against a real Redis instance since the in-memory fallback (which already refreshes unconditionally) can't reproduce this. bugbot caught this on review. --- .../hooks/model_based_tag_rate_limits_hook.py | 21 +++--- litellm/proxy/hooks/tag_rate_limits_shared.py | 23 ++++++- .../test_model_based_tag_rate_limits_hook.py | 67 +++++++++++++++---- 3 files changed, 88 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 2f1cf21f3af..e68dbc81208 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -934,12 +934,16 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] return built async def _check_and_increment_one( - self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int + self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool ) -> tuple[bool, float]: """Single-key atomic check-and-increment. Always one key per Lua - call -- see TAG_RL_CHECK_AND_INCR_SCRIPT's module docstring for why.""" + call -- see TAG_RL_CHECK_AND_INCR_SCRIPT's module docstring for why, + and for why `refresh_ttl` must be True for a concurrency key and + False for a requests key.""" if self._check_and_incr_script is not None: - raw: Final = await self._check_and_incr_script(keys=(key,), args=(limit, increment, ttl)) + raw: Final = await self._check_and_incr_script( + keys=(key,), args=(limit, increment, ttl, 1 if refresh_ttl else 0) + ) return bool(raw[0]), float(raw[1]) async with self._lock: @@ -962,7 +966,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] async def _atomic_check_and_increment( self, - checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], + checks: Sequence[tuple[InternalUsageCache, str, float, float, int, bool]], ) -> tuple[int | None, tuple[float, ...]]: """ All-or-nothing across every (cache, key, limit, increment, ttl) in @@ -1019,10 +1023,10 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # 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): + for index, (cache, key, limit, increment, ttl, refresh_ttl) in enumerate(checks): admitted = False try: - admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl) + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) finally: # Runs on a normal rejection (admitted stays False) and on # any exception/cancellation from the awaited call above @@ -1040,10 +1044,10 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] return None, tuple(admitted_values) async def _refund_admitted( - self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], up_to_index: int + self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int, bool]], up_to_index: int ) -> None: for refund_index in range(up_to_index): - refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index] + refund_cache, refund_key, _limit, refund_increment, _ttl, _refresh_ttl = checks[refund_index] try: await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment) except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest @@ -1150,6 +1154,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # with nothing to replace it. 0.0 if configured_limit.unit == "requests" and key in stale_request_keys else 1.0, self._ttl_for(configured_limit), + configured_limit.unit == "concurrency", ) for partition, (configured_limit, _tag_value, key) in zip(atomic_partitions, atomic_checks) ) diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index ad8e7bfe29d..f5a79f7de17 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -93,19 +93,36 @@ CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600 # `atomic_check_and_increment_by_n` in parallel_request_limiter_v3.py, applied # per-key instead of per-descriptor since each key already is one hash-tag # group by construction. +# +# refresh_ttl (ARGV[4]) distinguishes the two callers of this script: +# "requests" is an epoch-bucketed fixed window, whose TTL must be set once +# (at first write) and never extended, or the bucket outlives the epoch it's +# meant to reset at. "concurrency" is not windowed at all -- its TTL exists +# purely as a crash-safety net for a reservation whose explicit release never +# runs -- so a still-active bucket must keep pushing that TTL out on every +# admission, or a long-lived burst of continuous traffic expires the whole +# counter mid-flight (silently admitting past the cap, and letting a release +# for a since-reset counter decrement an unrelated, newer cohort). TAG_RL_CHECK_AND_INCR_SCRIPT: Final = """ local key = KEYS[1] local limit = tonumber(ARGV[1]) local increment = tonumber(ARGV[2]) local ttl = tonumber(ARGV[3]) +local refresh_ttl = tonumber(ARGV[4]) local current = tonumber(redis.call('GET', key) or 0) if current + increment > limit then return { 0, current } end local new_value = redis.call('INCRBY', key, increment) -local current_ttl = redis.call('TTL', key) -if current_ttl == -1 and ttl > 0 then - redis.call('EXPIRE', key, ttl) +if ttl > 0 then + if refresh_ttl == 1 then + redis.call('EXPIRE', key, ttl) + else + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl) + end + end end return { 1, new_value } """ 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 index 33edf7519a3..4a0f46e72aa 100644 --- 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 @@ -3340,6 +3340,49 @@ async def test_redis_backed_token_admission_sees_increments_the_in_memory_cache_ await redis_cache.async_delete_cache(key=token_key) +@pytest.mark.asyncio +async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + Bugbot finding: TAG_RL_CHECK_AND_INCR_SCRIPT only ran EXPIRE when a key + had no TTL at all, so a concurrency counter's expiry was fixed from its + first admission and never pushed out by later ones. A concurrency bucket + isn't epoch-windowed like requests/tokens/dollars -- its TTL exists only + as a crash-safety net for a reservation whose explicit release never + runs -- so a still-active bucket receiving continuous admissions must + keep extending that TTL, or it expires mid-flight under sustained + traffic, silently admitting past the cap. + """ + limiter, redis_cache = _redis_limiter(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + key = f"{{tag_rl:test:ttl-refresh:{uuid.uuid4().hex}}}:inflight" + cache = limiter.internal_usage_cache + try: + # A short, fixed ttl (bypassing _ttl_for's 3600s safety floor, which + # would make a real-time before/after comparison too slow to assert + # on deterministically) with refresh_ttl=True, matching how a + # concurrency check is actually admitted. + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = await redis_cache.redis_async_client.ttl(key) + assert ttl_after_first_admission > 0 + + await asyncio.sleep(2) + + # A second admission on the same still-live key, most of the way + # through the first admission's ttl, must push the ttl back out to + # the full window again, not leave it counting down toward zero. + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = await redis_cache.redis_async_client.ttl(key) + assert ttl_after_second_admission >= 2 + finally: + await redis_cache.async_delete_cache(key=key) + + # --------------------------------------------------------------------------- # team_public_model_name alias -- index lookup must not miss # --------------------------------------------------------------------------- @@ -3843,9 +3886,9 @@ async def test_refund_failure_on_one_key_does_not_block_others_or_raise(time_con 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), + (flaky.internal_usage_cache, failing_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, other_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, rejecting_key, 0.0, 1.0, 60, False), ] ) @@ -3870,18 +3913,18 @@ async def test_exception_mid_batch_refunds_every_earlier_admission_before_propag 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): + async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool): if key == raising_key: raise RuntimeError("simulated transient redis failure") - return await super()._check_and_increment_one(cache, key, limit, increment, ttl) + return await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now) with pytest.raises(RuntimeError): await flaky._atomic_check_and_increment( [ - (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60), - (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60), + (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60, False), ] ) @@ -3908,22 +3951,22 @@ async def test_a_raising_keys_own_ambiguous_outcome_is_never_refunded(time_contr 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): + async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool): if key == raising_key: # Simulate Redis committing the increment before the # response is lost: the write actually happens... - await super()._check_and_increment_one(cache, key, limit, increment, ttl) + await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) # ...but the caller never finds out. raise RuntimeError("simulated lost response after a committed redis write") - return await super()._check_and_increment_one(cache, key, limit, increment, ttl) + return await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now) with pytest.raises(RuntimeError): await flaky._atomic_check_and_increment( [ - (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60), - (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60), + (flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60, False), + (flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60, False), ] ) From 2ae037856a6bf6197d87b6c781158008e92bc200 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 16:26:11 -0400 Subject: [PATCH 17/44] fix(caching): apply each pipeline operation's own ttl in InMemoryCache async_increment_pipeline dropped each RedisPipelineIncrementOperation's own ttl field, so a counter created through it (Router's TPM/RPM tracking, parallel_request_limiter_v3's token/dollar accounting when Redis is absent, and this PR's own tag-based token/dollar limits) always fell back to the cache's 600-second default_ttl regardless of a real, often much longer, configured window. An hourly or daily limit's counter would silently expire and reset mid-window. allow_ttl_override already leaves a still-live ttl untouched on a later call, so threading the operation's ttl through on every increment only ever takes effect the first time. bugbot caught this on review. --- litellm/caching/in_memory_cache.py | 10 +++++- .../caching/test_in_memory_cache.py | 33 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 38a9966f9f9..7edf4672963 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -248,7 +248,15 @@ class InMemoryCache(BaseCache): ) -> list[float] | None: results: Final = [] for increment in increment_list: - result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) + # Each operation's own ttl must reach set_cache, or a key with no + # live ttl yet falls through to the cache's short default_ttl + # instead of the caller's real (often much longer) window -- + # allow_ttl_override already leaves an existing, still-live ttl + # untouched on a later increment, so passing this through on + # every call is safe: it only ever takes effect the first time. + result = await self.async_increment( + increment["key"], increment["increment_value"], ttl=increment.get("ttl"), **kwargs + ) results.append(result) return results diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 85e8308ae91..1e03159ca9a 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -45,6 +45,39 @@ async def test_async_increment_delegates_to_locked_sync_path(): assert cache.get_cache("counter") == 5 +async def test_async_increment_pipeline_applies_each_operations_own_ttl(): + """ + Bugbot finding: async_increment_pipeline dropped each operation's own + "ttl" field, so a counter created through it always got the cache's + 600-second default_ttl regardless of what the caller actually configured + (e.g. an hourly or daily rate-limit window) -- the counter (and whatever + it was tracking against a limit) silently reset mid-window. + """ + cache = InMemoryCache(default_ttl=600) + await cache.async_increment_pipeline([{"key": "long-window-counter", "increment_value": 1, "ttl": 7200}]) + + ttl_remaining = await cache.async_get_ttl("long-window-counter") + assert ttl_remaining is not None + assert ttl_remaining > time.time() + 600 + + +async def test_async_increment_pipeline_preserves_an_existing_live_ttl_on_later_increments(): + """ + A second increment on the same still-live counter must not reset its + remaining ttl back up to the full window -- only the first increment + (the one that actually creates the counter) should set it. + """ + cache = InMemoryCache(default_ttl=600) + await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}]) + ttl_after_first = cache.ttl_dict["counter"] + + await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}]) + ttl_after_second = cache.ttl_dict["counter"] + + assert ttl_after_second == ttl_after_first + assert cache.get_cache("counter") == 2 + + def test_in_memory_openai_obj_cache(): from openai import OpenAI From 0933fd9efb1d8723e2f714e7c8406663f0e01a2f Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 16:26:11 -0400 Subject: [PATCH 18/44] docs(proxy): document why a deployment-scoped breach rejects the whole hop bugbot flagged this as a bug (per-deployment-scoped limits should filter the over-limit deployment out of healthy_deployments and let a sibling serve, not reject the whole routing attempt). This is deliberate, documented design intent from the original plan: an earlier draft considered filter-and-retry semantics and rejected it, since rejecting the whole hop is simpler and avoids a caller silently succeeding against a deployment whose limit configuration they didn't intend to satisfy. Adding that reasoning as an inline comment so it doesn't get re-flagged as a bug on a future review. --- .../proxy/hooks/model_based_tag_rate_limits_hook.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index e68dbc81208..f292caf8cc3 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -739,6 +739,17 @@ def _classify_check( now: float, key_alias: str | None, ) -> _ClassifiedCheck | None: + # A breach of a deployment-scoped check (this function's own + # `deployment_scope is not None` branch, and every one of its callers' + # `_raise_over_limit`) deliberately rejects the whole routing attempt for + # this hop, not just the deployment(s) that own it -- it does not filter + # them out of `healthy_deployments` and let a sibling in the same group + # serve instead. An earlier design considered filter-and-retry-sibling + # semantics (matching how native tag routing filters candidates rather + # than rejecting the hop) and deliberately did not adopt it: rejecting + # the whole hop is simpler, and avoids a caller silently succeeding + # against a deployment whose limit configuration they didn't intend to + # satisfy. if configured_limit.deployment_scope is not None and not ( present_deployment_ids & frozenset(configured_limit.deployment_scope) ): From 9790a4ef0576728c0df44ed5814390c4b1c877ed Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 18:22:57 -0400 Subject: [PATCH 19/44] fix(proxy): mirror the concurrency ttl refresh onto the in-memory fallback, dedupe team-aliased buckets by internal model_name Two Bugbot findings from the same round: - The Redis refresh_ttl fix never reached the in-memory fallback path: async_set_cache called through unconditionally, and InMemoryCache's allow_ttl_override left a still-live ttl untouched regardless. Adds a refresh_ttl kwarg to InMemoryCache.set_cache/async_set_cache that bypasses that guard, wired through from the hook's own refresh_ttl flag. - A team-owned deployment resolved via its team_public_model_name alias got team_scope stamped into its bucket key, but the identical deployment resolved via its own internal model_name (which Router.should_include_deployment also permits for same-team or team-unconstrained callers) did not -- letting a caller split its usage across two independent counters by alternating which name it called with. Stamps the same team_scope onto the by_model_name entry whenever any deployment in that group has a team alias, so both paths resolve to the identical bucket. --- litellm/caching/in_memory_cache.py | 7 +++- .../hooks/model_based_tag_rate_limits_hook.py | 19 +++++++++- .../test_model_based_tag_rate_limits_hook.py | 37 +++++++++++++++++-- 3 files changed, 56 insertions(+), 7 deletions(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 7edf4672963..284e780cc6f 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -163,7 +163,12 @@ class InMemoryCache(BaseCache): return self.cache_dict[key] = value - if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl + # refresh_ttl bypasses allow_ttl_override's "leave a still-live ttl + # alone" guard -- a caller only sets it for a counter whose ttl must + # keep extending on every write (e.g. a concurrency reservation's + # crash-safety-net ttl), never for one that must stay fixed to its + # original epoch window (e.g. a fixed-period rate-limit bucket). + if kwargs.get("refresh_ttl") or self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index f292caf8cc3..1a7cb3f12f7 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -447,7 +447,20 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn 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 + model_name: ( + # `Router.should_include_deployment` lets a same-team caller + # reach a team-owned deployment by its own internal + # model_name, not only its team_public_model_name alias + # (litellm auto-generates a name unique per (team_id, uuid), + # so every deployment in this group shares one team_id when + # any does) -- stamping the identical team_scope here as the + # alias entry below gets keeps both paths resolving to the + # same bucket, so a caller can't split its usage across two + # independent counters just by alternating which name it calls. + tuple(replace(limit, team_scope=team_scope) for limit in configured) + if (team_scope := next((key[0] for dep in group if (key := _team_alias_key(dep))), None)) is not None + else configured + ) for model_name, deployment_group in groupby( sorted_by_model_name, key=lambda deployment: deployment["model_name"] ) @@ -963,7 +976,9 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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) + await cache.async_set_cache( + key=key, value=new_value, ttl=ttl, refresh_ttl=refresh_ttl, litellm_parent_otel_span=None + ) return True, new_value async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None: 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 index 4a0f46e72aa..ac24039758e 100644 --- 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 @@ -3383,6 +3383,32 @@ async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_co await redis_cache.async_delete_cache(key=key) +@pytest.mark.asyncio +async def test_in_memory_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + Bugbot finding: the Redis path's refresh_ttl fix above was never mirrored + onto the in-memory fallback, which called async_set_cache unconditionally + -- InMemoryCache.allow_ttl_override leaves a still-live ttl untouched, so + a concurrency counter's expiry stayed fixed from its first admission even + with refresh_ttl=True, the same silent-past-the-cap failure mode the + Redis fix closed. + """ + limiter = _make_limiter(time_controller) + cache = limiter.internal_usage_cache + in_memory_cache = cache.dual_cache.in_memory_cache + key = f"tag_rl:test:in-memory-ttl-refresh:{uuid.uuid4().hex}" + + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = in_memory_cache.ttl_dict[key] + + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = in_memory_cache.ttl_dict[key] + + assert ttl_after_second_admission > ttl_after_first_admission + + # --------------------------------------------------------------------------- # team_public_model_name alias -- index lookup must not miss # --------------------------------------------------------------------------- @@ -3396,6 +3422,12 @@ def test_build_limits_index_is_also_keyed_by_team_public_model_name(): model_group_alias). The index must resolve either name to the same configured limits, or a team-aliased chain's limits are silently never checked. + + Security regression: Router.should_include_deployment also lets a + same-team (or team-unconstrained) caller reach this deployment by its + own internal model_name, not only the alias. Both paths must resolve to + the identical team_scope, or a caller could split its usage across two + independent buckets just by alternating which name it calls with. """ deployment = _deployment( "real-model-name", @@ -3409,10 +3441,7 @@ def test_build_limits_index_is_also_keyed_by_team_public_model_name(): 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_name[0].team_scope == "team-1" assert by_alias[0].team_scope == "team-1" From 21ff11b13782e21b651ff11fa26b730901172483 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 07:15:19 -0400 Subject: [PATCH 20/44] fix(proxy): pin success accounting to admission's own routing-group snapshot resolve_any dedups a routing group's divergent per-deployment entries by picking the alphabetically first member model_name sharing a signature (resolved_group). Admission and success each independently rebuilt candidate_model_names from the router's live routing-group membership at their own point in time, so a deployment added or removed mid-request (a hot-reload) could make success pick a different resolved_group than admission did, hashing to a different Redis key and letting real token or dollar usage escape the bucket admission actually checked. Stashes admission's own candidate set on model_call_details, mirroring the existing admission-time-timestamp fix, so success reuses the identical snapshot. bugbot caught this on review. --- .../hooks/model_based_tag_rate_limits_hook.py | 53 +++++++++++++- .../test_model_based_tag_rate_limits_hook.py | 73 ++++++++++++++++++- 2 files changed, 121 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 1a7cb3f12f7..8ca8b39f0be 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -662,6 +662,21 @@ def _decode_reservations(raw: object) -> tuple[tuple[str, "_PartitionKey"], ...] # request, so its own most recent admission timestamp is the right one. _ADMISSION_TIME_FIELD: Final[str] = "_model_based_tag_rate_limits_admission_time" +# The routing-group membership (candidate_model_names) admission actually +# resolved against, stashed the same way _ADMISSION_TIME_FIELD is. resolve_any +# dedupes divergent per-deployment entries by picking the alphabetically first +# member model_name sharing a signature (resolved_group) -- a pure function of +# this exact candidate set. Router's live routing-group membership can change +# between admission and success (a deployment added or removed mid-request via +# /model/new or a config hot-reload), and success independently re-deriving +# candidate_model_names from *live* membership at that later point can pick a +# different resolved_group than admission did, hashing to a different Redis +# key -- so success accounting silently misses the bucket admission actually +# checked, letting real usage escape the enforced cap. Reusing admission's own +# snapshot keeps resolve_any's output identical at both points regardless of +# what changed in between. +_ROUTING_GROUP_CANDIDATES_FIELD: Final[str] = "_model_based_tag_rate_limits_routing_group_candidates" + class _TagRateLimitIndex: """Rebuilds the limits index when `llm_router.model_list` changes, or at @@ -878,6 +893,25 @@ def _admission_time_or(kwargs: Mapping[str, object], fallback: float) -> float: return recorded if isinstance(recorded, float) else fallback +def _record_routing_group_candidates( + request_kwargs: Mapping[str, object], candidate_model_names: tuple[str, ...] +) -> None: + """Stash the routing-group membership admission resolved against -- see + `_ROUTING_GROUP_CANDIDATES_FIELD`'s docstring for why. Silently a no-op + without a real logging object (defensive only; every real request has + one): success accounting falls back to its own live reconstruction, same + as before this fix existed.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if isinstance(model_call_details, dict): + model_call_details[_ROUTING_GROUP_CANDIDATES_FIELD] = candidate_model_names + + +def _routing_group_candidates_or(kwargs: Mapping[str, object], fallback: tuple[str, ...]) -> tuple[str, ...]: + recorded: Final = kwargs.get(_ROUTING_GROUP_CANDIDATES_FIELD) + return recorded if isinstance(recorded, tuple) else fallback + + @dataclass(frozen=True, slots=True) class _CachePartition: internal_usage_cache: InternalUsageCache @@ -1106,9 +1140,11 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # 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. + # below). Deriving both sides from the same full-membership source + # handles cooldown-state drift between the two points in time; actual + # membership drift (a deployment added or removed mid-request) still + # needs admission's own snapshot stashed and reused -- see + # _ROUTING_GROUP_CANDIDATES_FIELD's docstring. routing_group_deployments: Final = self.llm_router._get_routing_group_deployments( # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching resolve_any's own reliance on this method model=model, team_id=team_id ) @@ -1117,6 +1153,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if routing_group_deployments is not None else tuple(name for d in healthy_deployments if isinstance(name := d.get("model_name"), str)) ) + _record_routing_group_candidates(resolved_request_kwargs, candidate_model_names) configured: Final = self._index.get(self.llm_router).resolve_any(model, team_id, candidate_model_names) if not configured: return healthy_deployments @@ -1621,6 +1658,13 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # only when `model_group` isn't a routing group at all (a plain # single-model_name chain, where resolve() already matches directly # and this candidate set is never actually consulted). + # + # Reconstructing live membership here is itself only a fallback: it + # can still disagree with admission's own candidate set if the + # routing group's actual membership changed between the two points + # in time (not just cooldown/health state) -- _routing_group_candidates_or + # below prefers admission's own stashed snapshot whenever one exists. + # See _ROUTING_GROUP_CANDIDATES_FIELD's docstring. deployment_id: Final = standard_logging_object.get("model_id") serving_deployment: Final = ( self.llm_router.get_deployment(deployment_id) if isinstance(deployment_id, str) else None @@ -1628,11 +1672,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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 = ( + live_candidate_model_names: Final = ( tuple(dep["model_name"] for dep in routing_group_deployments) if routing_group_deployments is not None else ((serving_deployment.model_name,) if serving_deployment is not None else ()) ) + candidate_model_names: Final = _routing_group_candidates_or(kwargs, fallback=live_candidate_model_names) configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names) if not configured: return 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 index ac24039758e..a3c9e80e502 100644 --- 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 @@ -33,7 +33,7 @@ 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 +from litellm.types.router import Deployment, RoutingGroup, TagRateLimitEntry, TagRateLimitScope class TimeController: @@ -1596,6 +1596,77 @@ async def test_log_success_event_accounts_against_the_same_bucket_admission_chec ) +@pytest.mark.asyncio +async def test_log_success_event_uses_admissions_own_candidate_set_when_group_membership_drifts(time_controller): + """ + Bugbot finding: resolve_any's dedup picks the alphabetically first member + model_name sharing a signature as resolved_group, a pure function of + candidate_model_names. Both admission and success independently rebuild + that set from the router's *live* routing-group membership, so a + deployment added mid-request (a hot-reload) whose name sorts earlier can + make success pick a different resolved_group than admission did, + accounting real usage into a bucket admission never checked. + """ + token_limits = { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + } + router = litellm.Router( + model_list=[ + _deployment("backend-a", "dep-a", token_limits), + _deployment("backend-b", "dep-b", token_limits), + ], + routing_groups=[ + RoutingGroup(group_name="my-group", models=["backend-a", "backend-b"], routing_strategy="simple-shuffle") + ], + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + request_kwargs, model_call_details = _call_context(["end_user_id:u1"]) + + healthy = router._get_routing_group_deployments(model="my-group", team_id=None) + admitted = await limiter.async_filter_deployments( + model="my-group", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert admitted == healthy + admission_bucket_group = limiter._index.get(router).resolve_any( + "my-group", team_id=None, candidate_model_names=("backend-a", "backend-b") + )[0].resolved_group + + router.get_routing_group("my-group").models.append("backend-0") + router.add_deployment( + Deployment( + model_name="backend-0", + litellm_params={"model": "gpt-4o", "mock_response": "ok"}, # type: ignore + model_info={"id": "dep-0", "tag_rate_limits": token_limits}, + ) + ) + serving_deployment_id = "dep-b" if admission_bucket_group == "backend-a" else "dep-a" + + model_call_details["standard_logging_object"] = { + "model_group": "my-group", + "model_id": serving_deployment_id, + "total_tokens": 42, + "response_cost": 0.01, + } + await limiter.async_log_success_event(kwargs=model_call_details, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + admission_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group=admission_bucket_group, limit=500000 + ) + drifted_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-0", limit=500000 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=admission_key, litellm_parent_otel_span=None)) + == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=drifted_key, litellm_parent_otel_span=None) is None + + @pytest.mark.asyncio async def test_admission_dedups_against_the_full_group_not_just_currently_healthy_members(time_controller): """ From b98195e3a40cfa59842cb1284898a0c1e624848b Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 07:44:27 -0400 Subject: [PATCH 21/44] fix(proxy): use typed Deployment/ModelInfo construction in the new drift regression test LiteLLM_Params and ModelInfo instead of raw dicts, plus assert narrowing on two Router calls that can return None, to satisfy basedpyright's strict Deployment(...) construction path. --- .../hooks/test_model_based_tag_rate_limits_hook.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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 index a3c9e80e502..6fcaf865c5d 100644 --- 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 @@ -33,7 +33,7 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( BACKGROUND_TASKS as _BACKGROUND_TASKS, CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, ) -from litellm.types.router import Deployment, RoutingGroup, TagRateLimitEntry, TagRateLimitScope +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, RoutingGroup, TagRateLimitEntry, TagRateLimitScope class TimeController: @@ -1626,6 +1626,7 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me request_kwargs, model_call_details = _call_context(["end_user_id:u1"]) healthy = router._get_routing_group_deployments(model="my-group", team_id=None) + assert healthy is not None admitted = await limiter.async_filter_deployments( model="my-group", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs ) @@ -1634,12 +1635,14 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me "my-group", team_id=None, candidate_model_names=("backend-a", "backend-b") )[0].resolved_group - router.get_routing_group("my-group").models.append("backend-0") + routing_group = router.get_routing_group("my-group") + assert routing_group is not None + routing_group.models.append("backend-0") router.add_deployment( Deployment( model_name="backend-0", - litellm_params={"model": "gpt-4o", "mock_response": "ok"}, # type: ignore - model_info={"id": "dep-0", "tag_rate_limits": token_limits}, + litellm_params=LiteLLM_Params(model="gpt-4o", mock_response="ok"), + model_info=ModelInfo(id="dep-0", tag_rate_limits=token_limits), ) ) serving_deployment_id = "dep-b" if admission_bucket_group == "backend-a" else "dep-a" From 563b3d6adf0ff7480a5e90fb30aca43ef5d850be Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 12:15:19 -0400 Subject: [PATCH 22/44] fix(proxy): resolve pre-existing basedpyright reportArgumentType/reportCallIssue errors model_based_tag_rate_limits_hook.py had 3 real type errors present since its introduction, invisible in scoped per-file checks but visible once codebase-wide upstream drift finally pushed the totals over ceiling: sorted()/groupby() keyed by a raw dict lookup returning object rather than a provably orderable type, and a tuple passed where async_increment_tokens_with_ttl_preservation expects a list. Adds a small typed _model_name_of accessor for the first two, and drops an unneeded tuple() conversion for the third since the value was already a list. Ran make lint-budget-update per repo convention. --- .../hooks/model_based_tag_rate_limits_hook.py | 17 ++++++++++++----- type-discipline-budget.json | 2 +- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 8ca8b39f0be..79005bf71b6 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -175,6 +175,15 @@ def _deployment_id(deployment: Mapping[str, object]) -> str | None: return (deployment.get("model_info") or _EMPTY_MAPPING).get("id") +def _model_name_of(deployment: Mapping[str, object]) -> str: + """`model_name` is a required field on every deployment dict; `str(...)` + (rather than a bare index) gives `sorted`/`groupby`'s key functions a + provably orderable, hashable return type without an unsafe cast -- every + real deployment's value here is already a string, so this is a no-op + coercion in practice.""" + return str(deployment["model_name"]) + + def _extract_team_id(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: """Reads `user_api_key_team_id` from only the one field `get_metadata_variable_name_from_kwargs` names as authoritative for this @@ -444,7 +453,7 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn 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"]) + sorted_by_model_name: Final = sorted(model_list, key=_model_name_of) by_model_name: Final[Mapping[str, tuple[_ConfiguredLimit, ...]]] = MappingProxyType( { model_name: ( @@ -461,9 +470,7 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn if (team_scope := next((key[0] for dep in group if (key := _team_alias_key(dep))), None)) is not None else configured ) - for model_name, deployment_group in groupby( - sorted_by_model_name, key=lambda deployment: deployment["model_name"] - ) + for model_name, deployment_group in groupby(sorted_by_model_name, key=_model_name_of) for group in (tuple(deployment_group),) if (configured := tuple(limit for unit in _LIMIT_UNITS for limit in _build_group_limits(group, unit))) } @@ -1727,7 +1734,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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), + pipeline_operations=group_operations, parent_otel_span=parent_otel_span, ) ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 959b2eada25..644c0f22063 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16616 + "limit": 16615 }, "LIT011": { "limit": 5583 From bd673268d3c5c6486a1b5a6acc584f5087bdb11b Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 17:13:07 -0400 Subject: [PATCH 23/44] fix(proxy): unify alias/internal-name tag rate-limit buckets and close caller-forged metadata bypass Bugbot finding: stamping team_scope on both the team-alias and internal model_name paths didn't unify their Redis keys, since _hash_tag still hashed the caller-visible name by default and that name differs between the two paths. Both index branches now also stamp resolved_group with the team's own public alias, so a team calling its own internal model_name and the same team calling its public alias land on one shared bucket. Veria AI finding: admission resolved the authoritative metadata bucket via get_metadata_variable_name_from_kwargs, which only checks key presence. A caller could forge an empty (or None) litellm_metadata to make admission read no tags/team at all and bypass every configured limit. Admission now reuses the same truthiness-checking resolver already used for success accounting, renamed to reflect both call sites. --- .../hooks/model_based_tag_rate_limits_hook.py | 102 +++++++++----- litellm/proxy/hooks/tag_rate_limits_shared.py | 29 ++-- .../test_model_based_tag_rate_limits_hook.py | 133 ++++++++++++++++++ 3 files changed, 214 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 79005bf71b6..45ae5943598 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -27,7 +27,6 @@ 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 @@ -99,7 +98,7 @@ 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, + resolve_authoritative_metadata_variable_name as _resolve_authoritative_metadata_variable_name, ) from litellm.proxy.hooks.tag_rate_limits_shared import ( scope_signature as _scope_signature, @@ -150,24 +149,33 @@ class _ConfiguredLimit: # 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). + # The owning team_id, set whenever this limit belongs to a team-owned + # deployment -- via `by_team_alias`, or via `by_model_name` for a group + # any of whose deployments also declare a `team_public_model_name` (see + # `_build_limits_index`). team_public_model_name is only unique per team, + # so two teams can publish the identical alias string; without the + # team_id folded into the bucket key too, both teams' identically-named, + # identically-configured limits would collide on the same Redis counter + # despite the index itself correctly scoping the `by_team_alias` lookup + # by (team_id, alias). team_scope: str | None = None - # 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. + # Overrides the caller-visible name in the bucket key whenever that name + # doesn't uniquely identify the bucket. Two independent cases set this: + # (1) `resolve()`'s direct lookup missed and `resolve_any()` fell back to + # resolving via a candidate deployment's own model_name instead (routing + # groups, and any other indirection where Router deliberately keeps the + # caller-visible name distinct from every deployment's own model_name) -- + # here it's stamped with that deployment's own model_name, so limits from + # two different underlying model_names sharing one routing group never + # collide on one counter; (2) the limit belongs to a team-owned + # deployment (`team_scope` is set) -- here it's stamped with the team's + # own `team_public_model_name`, the same value regardless of whether this + # entry was found via `by_team_alias` (caller used the public alias) or + # via `by_model_name` (caller used the deployment's own internal, + # auto-generated model_name), so both call shapes land on one shared + # bucket instead of splitting a team's usage across two counters. None + # when resolved directly and not team-owned, in which case the + # caller-visible name is already unambiguous and safe to hash by. resolved_group: str | None = None @@ -186,8 +194,8 @@ def _model_name_of(deployment: Mapping[str, object]) -> str: 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 + `_resolve_authoritative_metadata_variable_name` names as authoritative for + this request -- never falling back to the other field, since `litellm_pre_call_utils.py` writes the real, server-authenticated value into that one field alone and leaves the other exactly as the caller sent it. An OR-fallback across both would let a caller's own @@ -462,12 +470,16 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn # model_name, not only its team_public_model_name alias # (litellm auto-generates a name unique per (team_id, uuid), # so every deployment in this group shares one team_id when - # any does) -- stamping the identical team_scope here as the - # alias entry below gets keeps both paths resolving to the - # same bucket, so a caller can't split its usage across two + # any does). Stamping team_scope alone is not enough to unify + # this with the alias entry below: `_hash_tag` still hashes + # the caller-visible name by default, and that name differs + # between the two paths (the internal model_name here vs. the + # public alias below). Also stamping `resolved_group` with the + # team's own alias forces both paths to hash under the + # identical name, so a caller can't split its usage across two # independent counters just by alternating which name it calls. - tuple(replace(limit, team_scope=team_scope) for limit in configured) - if (team_scope := next((key[0] for dep in group if (key := _team_alias_key(dep))), None)) is not None + tuple(replace(limit, team_scope=team_key[0], resolved_group=team_key[1]) for limit in configured) + if (team_key := next((key for dep in group if (key := _team_alias_key(dep))), None)) is not None else configured ) for model_name, deployment_group in groupby(sorted_by_model_name, key=_model_name_of) @@ -487,7 +499,13 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn for aliased_group in (tuple(dep for _key, dep in alias_group),) if ( alias_configured := tuple( - replace(limit, team_scope=alias_key[0]) + # resolved_group is already the caller-visible name on + # this path (the caller reached this group by dialing the + # alias directly), but stamping it explicitly keeps both + # index branches symmetric and independent of whatever + # value the caller happens to pass as `model_group` into + # `_hash_tag`. + replace(limit, team_scope=alias_key[0], resolved_group=alias_key[1]) for unit in _LIMIT_UNITS for limit in _build_group_limits(aliased_group, unit) ) @@ -711,12 +729,15 @@ def _scope_suffix(deployment_scope: tuple[str, ...] | None) -> str: 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. + # resolved_group overrides the caller-visible model_group in two cases: + # resolve_any()'s per-deployment fallback (routing groups), where the + # caller-visible name is ambiguous (shared by every member model_name), so + # hashing by it would collide two different underlying model_names' + # identically-named limits onto one counter; and a team-owned deployment, + # where the caller-visible name differs depending on whether the caller + # dialed the team's public alias or the deployment's own internal + # model_name, so hashing by it would split one team's usage across two + # counters. See _ConfiguredLimit.resolved_group. effective_model_group: Final = configured.resolved_group if configured.resolved_group is not None else model_group scope: Final = _scope_suffix(configured.deployment_scope) # team_scope disambiguates two teams that publish the identical @@ -1139,7 +1160,14 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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) + # Not `get_metadata_variable_name_from_kwargs` (naive key-presence + # check): a caller can forge an empty (or `None`) `litellm_metadata` + # on an ordinary request to make that check pick it over the real, + # populated `metadata` the proxy wrote authenticated team/tag identity + # into, seeing no tags at all and admitting past every configured + # limit. See `_resolve_authoritative_metadata_variable_name`'s own + # docstring. + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(resolved_request_kwargs) team_id: Final = _extract_team_id(resolved_request_kwargs, metadata_variable_name) # Built from the full routing-group membership, not `healthy_deployments` # (Router's own cooldown-filtered list for this hop): a member that's @@ -1499,12 +1527,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # 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. + # `_resolve_authoritative_metadata_variable_name`'s own docstring. litellm_params_raw: Final = kwargs.get("litellm_params") litellm_params_for_metadata: Final = ( litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs ) - metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata) + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) try: await self.internal_usage_cache.dual_cache.async_delete_cache( @@ -1644,7 +1672,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # 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) + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) team_id: Final = _extract_team_id(litellm_params_for_metadata, metadata_variable_name) key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name) diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index f5a79f7de17..92c228ad819 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -221,21 +221,24 @@ def entry_applies(entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str return key_alias in entry.apply_to_key_alias -def resolve_success_event_metadata_variable_name( - litellm_params_for_metadata: Mapping[str, object], +def resolve_authoritative_metadata_variable_name( + metadata_source: Mapping[str, object], ) -> Literal["metadata", "litellm_metadata"]: """`get_metadata_variable_name_from_kwargs` only checks key presence, which - misresolves at `async_log_success_event` time: `kwargs["litellm_params"]` - always carries a `litellm_metadata` key (typically `None`) alongside the - real, populated `metadata` dict for a standard (non - LITELLM_METADATA_ROUTES) request, so the key-presence check always picks - `litellm_metadata` there and silently reads no tags/identity at all. - Requiring the value to actually be a populated dict, matching - `_get_request_tags`'s own truthiness check in litellm_logging.py, only - ever prefers `litellm_metadata` when it is genuinely the field the proxy - wrote identity/tags into (LITELLM_METADATA_ROUTES pre-seed it before - admission runs, so it is always a populated dict by success time there).""" - litellm_metadata: Final = litellm_params_for_metadata.get("litellm_metadata") + misresolves both at admission time and at `async_log_success_event` time: + a caller can forge an empty (or merely present-but-`None`) `litellm_metadata` + on an ordinary request -- `kwargs["litellm_params"]` also always carries a + `litellm_metadata` key (typically `None`) alongside the real, populated + `metadata` dict for a standard (non LITELLM_METADATA_ROUTES) request -- and + the key-presence check always picks `litellm_metadata` in both cases, + silently reading no tags/identity at all and admitting the request against + every configured limit. Requiring the value to actually be a populated + dict, matching `_get_request_tags`'s own truthiness check in + litellm_logging.py, only ever prefers `litellm_metadata` when it is + genuinely the field the proxy wrote identity/tags into + (LITELLM_METADATA_ROUTES pre-seed it before admission runs, so it is + always a populated dict there, both at admission and at success time).""" + litellm_metadata: Final = metadata_source.get("litellm_metadata") if isinstance(litellm_metadata, Mapping) and litellm_metadata: return "litellm_metadata" return "metadata" 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 index 6fcaf865c5d..e9288566da0 100644 --- 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 @@ -214,6 +214,66 @@ def test_extract_team_id_ignores_a_forged_value_in_the_non_authoritative_field() assert _extract_team_id(request_kwargs, "litellm_metadata") == "real-team" +@pytest.mark.asyncio +async def test_filter_deployments_ignores_a_forged_empty_litellm_metadata_key(time_controller): + """ + Veria AI finding: get_metadata_variable_name_from_kwargs picks + "litellm_metadata" whenever that key is merely present, regardless of its + value. add_litellm_data_to_request writes real, authenticated team/tag + identity into "metadata" for an ordinary (non LITELLM_METADATA_ROUTES) + request, but leaves any caller-supplied "litellm_metadata" sitting + alongside it untouched -- so a caller adding an empty "litellm_metadata" + to a chat completions request made admission read no tags/team at all, + sailing past every configured limit. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": {}} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + +@pytest.mark.asyncio +async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_present_but_none(time_controller): + """ + Same misresolution, naturally occurring rather than attacker-forged: a + plain chat completion's kwargs carries a "litellm_metadata" key that is + always present but set to None, alongside the real, populated "metadata" + dict (see resolve_authoritative_metadata_variable_name's own docstring). + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": None} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # --------------------------------------------------------------------------- # TagRateLimitEntry -- limit validation # --------------------------------------------------------------------------- @@ -3519,6 +3579,33 @@ def test_build_limits_index_is_also_keyed_by_team_public_model_name(): assert by_alias[0].team_scope == "team-1" +def test_build_limits_index_computes_identical_bucket_key_for_alias_and_internal_model_name(): + """ + Bugbot finding: matching team_scope alone does not unify the two paths' + buckets, because _hash_tag hashes the caller-visible model_group by + default, and that name is `real-model-name` on the by_model_name path but + `team-alias-name` on the by_team_alias path. Both entries must also share + the identical resolved_group (the team's own alias) so _hash_tag hashes + both under one name -- otherwise a team calling its own internal + model_name lands on a different Redis counter than the same team calling + its public alias, doubling its effective quota. + """ + deployment = _deployment( + "real-model-name", + "dep-1", + {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}}, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + index = _build_limits_index([deployment]) + by_name = index.resolve("real-model-name", team_id=None)[0] + by_alias = index.resolve("team-alias-name", team_id="team-1")[0] + + key_via_internal_name = _bucket_key("real-model-name", by_name, tag_value="u1", bucket_id=0) + key_via_alias = _bucket_key("team-alias-name", by_alias, tag_value="u1", bucket_id=0) + assert key_via_internal_name == key_via_alias + + def test_build_limits_index_preserves_key_ttl_seconds_and_max_in_memory_cache_size(): """ Regression test: _configured_limit_for_signature used to reconstruct a @@ -3707,6 +3794,52 @@ async def test_filter_deployments_enforces_limit_when_called_with_team_alias(tim ) +@pytest.mark.asyncio +async def test_filter_deployments_shares_quota_across_alias_and_internal_model_name(time_controller): + """ + Bugbot finding on a prior fix for this same drift: stamping the identical + team_scope on both paths was not enough, since _hash_tag still hashes the + caller-visible name (the alias here, the deployment's own internal + model_name there) by default. A team calling through its public alias and + the identical team calling through the deployment's own internal + model_name must draw from the same bucket, or the team gets one quota per + name it happens to call with -- doubling its real limit. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "real-model-name", + "dep-1", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # First call exhausts the limit of 1 via the team's public alias. + await limiter.async_filter_deployments( + model="team-alias-name", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}}, + ) + # Router.should_include_deployment also lets the same team reach this + # deployment by its own internal model_name; that call must be rejected + # against the alias call's own bucket, not admitted into a fresh one. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="real-model-name", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}}, + ) + + @pytest.mark.asyncio async def test_filter_deployments_does_not_cross_team_alias_boundary(time_controller): """ From edc548af97046f8604964549067cdf11982733dd Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 17:58:18 -0400 Subject: [PATCH 24/44] fix(proxy): require an unforgeable marker to trust litellm_metadata, and stop admission from releasing a live sibling's concurrency reservation Bugbot finding: resolve_authoritative_metadata_variable_name treated any non-empty litellm_metadata as authoritative, but a caller can populate it with unrelated keys on an ordinary route where "metadata" is the real field. Now requires the unconditionally-stamped, strip-protected "user_api_key_auth" marker instead of mere truthiness. Veria AI finding: Router.abatch_completion's comma-separated multi-model dispatch runs branches concurrently as separate asyncio Tasks that all share one litellm_logging_obj, so a still-live sibling branch's own concurrency reservation could sit in the same model_call_details a new hop's admission was cleaning up. _release_stale_hop_reservations now only reclaims entries queued by its own asyncio.Task, leaving a differently tasked (still-live) entry alone. --- .../hooks/model_based_tag_rate_limits_hook.py | 81 ++++++-- litellm/proxy/hooks/tag_rate_limits_shared.py | 19 +- .../test_model_based_tag_rate_limits_hook.py | 177 ++++++++++++++++-- 3 files changed, 239 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 45ae5943598..4c32f822b42 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -556,6 +556,15 @@ _INDEX_TTL_SECONDS: Final = 5.0 # server-side per logical request (and shared across that request's own # fallback hops, matching the original chain-wide release semantics), so it # can't be forged or guessed. +# +# "shared across that request's own fallback hops" is not the same as +# "scoped to one asyncio Task": `Router.abatch_completion`'s comma-separated +# multi-model dispatch runs several branches concurrently, each its own +# Task, but hands every branch the identical `litellm_logging_obj` -- so +# each entry also carries the Task that queued it (see +# `_queue_pending_reservations`), letting `_release_stale_hop_reservations` +# tell a genuinely stale same-task hop apart from a still-live sibling +# branch's own reservation. _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 @@ -893,6 +902,17 @@ def _queue_pending_reservations( logging object (defensive only; every real request has one): a queued concurrency reservation still self-heals via `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, just later. + + Each entry is stamped with the queueing coroutine's own `asyncio.Task`: + `Router.abatch_completion`'s comma-separated multi-model dispatch runs + several `acompletion` calls concurrently as *separate* tasks that all + share one `model_call_details` (the proxy attaches one `litellm_logging_obj` + to the request before the comma-split, and every branch inherits that + same reference), so this field is no longer scoped to one logical + request's own serial fallback chain the way its docstring assumes. + `_release_stale_hop_reservations` uses the stamp to tell "an earlier hop + of *this* chain, safe to reclaim" apart from "a concurrent sibling + branch's own still-live reservation," which must never be touched here. """ logging_obj: Final = request_kwargs.get("litellm_logging_obj") model_call_details: Final = getattr(logging_obj, "model_call_details", None) @@ -902,7 +922,10 @@ def _queue_pending_reservations( 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 + current_task: Final = asyncio.current_task() + pending.extend( + (key, partition_key, current_task) for key, partition_key in reservations + ) # mutable-ok: see comment above def _record_admission_time(request_kwargs: Mapping[str, object], now: float) -> None: @@ -1460,13 +1483,25 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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. + runs, *within the same asyncio Task*, can only belong to an earlier + hop of this same request's own fallback chain that already concluded + and failed: Router awaits one hop's entire attempt (call plus its own + failure handling) before starting the next, and a hop that instead + succeeded ends the request there via async_log_success_event, which + already pops everything -- so admission is never re-entered, in that + same task, while an earlier hop's reservation is still legitimately + in flight. + + The task check matters because `model_call_details` is not always + scoped to one such chain: `Router.abatch_completion`'s comma-separated + multi-model dispatch runs several branches concurrently, each its own + Task, but every branch shares the identical `litellm_logging_obj` (see + `_queue_pending_reservations`'s own docstring) -- so a reservation + queued by a still-running sibling branch can be sitting here too, and + releasing it out from under that branch would let more calls through + a concurrency limit than it allows. Only entries this exact Task + queued are safe to treat as stale; anything else is left for its own + branch to release. LiteLLM only invokes a request's CustomLogger.async_log_failure_event once per request, for whichever hop fails first (its internal @@ -1506,16 +1541,17 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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) + release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_current_task=True) if release_keys: await self._release_keys(release_keys) pending_request_increments: Final = model_call_details.get(_PENDING_REQUEST_INCREMENTS_FIELD) if not isinstance(pending_request_increments, list): return frozenset() - return frozenset(key for key, _partition_key in pending_request_increments) + current_task: Final = asyncio.current_task() + return frozenset(key for key, _partition_key, task in pending_request_increments if task is current_task) async def _pop_pending_concurrency_keys( - self, kwargs: Mapping[str, object] + self, kwargs: Mapping[str, object], *, only_current_task: bool = False ) -> 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 @@ -1544,21 +1580,26 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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. + # Snapshot then remove only those exact entries, never a blanket + # clear: a sibling branch sharing this same request's + # model_call_details can still be live and appending concurrently + # (see the field's own docstring), so wiping the whole list here + # would silently strand that branch's reservation instead of + # releasing it later. `only_current_task` additionally excludes any + # entry a *different*, still-running Task queued -- see + # `_release_stale_hop_reservations`'s own docstring for why that + # distinction, not just presence, decides what's actually stale. pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) if not isinstance(pending, list) or not pending: return () - keys: Final = tuple(pending) - for key in keys: + current_task: Final = asyncio.current_task() + snapshot: Final = tuple(entry for entry in pending if not only_current_task or entry[2] is current_task) + for entry in snapshot: try: - pending.remove(key) + pending.remove(entry) except ValueError: pass - return keys + return tuple(entry[:2] for entry in snapshot) async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: """ diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index 92c228ad819..dafc1b0a95d 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -232,14 +232,19 @@ def resolve_authoritative_metadata_variable_name( `metadata` dict for a standard (non LITELLM_METADATA_ROUTES) request -- and the key-presence check always picks `litellm_metadata` in both cases, silently reading no tags/identity at all and admitting the request against - every configured limit. Requiring the value to actually be a populated - dict, matching `_get_request_tags`'s own truthiness check in - litellm_logging.py, only ever prefers `litellm_metadata` when it is - genuinely the field the proxy wrote identity/tags into - (LITELLM_METADATA_ROUTES pre-seed it before admission runs, so it is - always a populated dict there, both at admission and at success time).""" + every configured limit. + + Merely requiring the value to be a non-empty dict is not enough either: a + caller can populate its own, unrelated keys on the non-authoritative + bucket (e.g. `{"litellm_metadata": {"x": 1}}` on an ordinary route), which + is non-empty but still not the field the proxy wrote identity into. + `add_litellm_data_to_request` unconditionally stamps `user_api_key_auth` + into whichever bucket the route actually resolved as authoritative, and + strips any `user_api_key_`-prefixed key a caller pre-populates on the + other bucket -- so requiring that marker's presence, not mere + truthiness, can't be forged onto the wrong side.""" litellm_metadata: Final = metadata_source.get("litellm_metadata") - if isinstance(litellm_metadata, Mapping) and litellm_metadata: + if isinstance(litellm_metadata, Mapping) and "user_api_key_auth" in litellm_metadata: return "litellm_metadata" return "metadata" 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 index e9288566da0..60df49728f1 100644 --- 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 @@ -274,6 +274,39 @@ async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_presen ) +@pytest.mark.asyncio +async def test_filter_deployments_ignores_a_forged_populated_litellm_metadata_key(time_controller): + """ + Bugbot finding on the fix above: requiring litellm_metadata to be merely + non-empty is still forgeable -- a caller can populate it with its own, + unrelated keys on an ordinary route, which is non-empty but carries none + of the real identity the proxy wrote into "metadata". Only + add_litellm_data_to_request's own "user_api_key_auth" marker, stripped + from any bucket the caller doesn't own, proves a bucket is authoritative. + """ + limiter = _make_limiter(time_controller) + deployment = _deployment( + "grp", + "dep-1", + {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + ) + router = litellm.Router(model_list=[deployment]) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + request_kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "litellm_metadata": {"x": 1}, + } + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # --------------------------------------------------------------------------- # TagRateLimitEntry -- limit validation # --------------------------------------------------------------------------- @@ -1438,6 +1471,54 @@ async def test_log_success_event_accounts_when_litellm_params_carries_a_null_lit ) +@pytest.mark.asyncio +async def test_log_success_event_ignores_a_forged_populated_litellm_metadata_key(time_controller): + """ + Same misresolution as the null-key case above, but with a populated + (not merely present) forged litellm_metadata -- a caller-supplied dict + with unrelated keys is still not the field the proxy wrote identity + into. Only the unconditionally-stamped "user_api_key_auth" marker, + which litellm_pre_call_utils.py strips from any bucket a caller doesn't + own, proves a bucket is authoritative. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "litellm_params": { + "litellm_metadata": {"x": 1}, + "metadata": {"tags": ["end_user_id:u1"]}, + }, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + @pytest.mark.asyncio async def test_log_success_event_accounts_the_key_backed_tag_not_a_caller_forged_one(time_controller): """ @@ -1523,7 +1604,7 @@ async def test_log_success_event_reads_nested_litellm_metadata_when_that_is_auth kwargs = { "litellm_params": { "metadata": {"tags": []}, - "litellm_metadata": {"tags": ["end_user_id:u1"]}, + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_auth": {}}, }, "standard_logging_object": { "model_group": "grp", @@ -1923,7 +2004,9 @@ async def test_log_success_event_accounts_against_the_team_id_admission_checked( # 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"}} + request_kwargs = { + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1", "user_api_key_auth": {}} + } result = await limiter.async_filter_deployments( model="team-alias-name", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs ) @@ -1933,7 +2016,9 @@ async def test_log_success_event_accounts_against_the_team_id_admission_checked( # 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"}}, + "litellm_params": { + "litellm_metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1", "user_api_key_auth": {}} + }, "standard_logging_object": { "model_group": "team-alias-name", "model_id": "dep-1", @@ -2831,6 +2916,48 @@ async def test_next_hops_admission_releases_a_prior_hops_leaked_reservation(time assert result == healthy +@pytest.mark.asyncio +async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_controller): + """ + Veria AI finding: Router.abatch_completion's comma-separated multi-model + dispatch runs each model concurrently as its own asyncio.Task, but every + branch is handed the identical litellm_logging_obj (the proxy attaches + one to the request before the comma-split), so two genuinely concurrent + branches share one model_call_details. Before the fix, a second branch's + admission-time stale-reservation cleanup couldn't tell that apart from + an earlier, already-failed hop of its own retry chain, so it released + the first branch's still-live reservation and let both branches through + a concurrency limit of 1. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + first_admitted = asyncio.Event() + + async def _branch_one() -> None: + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + first_admitted.set() + # Held "in flight" while the second branch's admission runs, exactly + # like two concurrently in-flight provider calls. + await asyncio.sleep(0.05) + + async def _branch_two() -> None: + await first_admitted.wait() + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + results = await asyncio.gather( + asyncio.create_task(_branch_one()), asyncio.create_task(_branch_two()), return_exceptions=True + ) + rejections = [result for result in results if isinstance(result, ProxyRateLimitError)] + assert len(rejections) == 1 + + def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[ @@ -3985,13 +4112,16 @@ def test_concurrency_ttl_floor_does_not_shorten_a_longer_period_seconds(): @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"]} + # Entries are (key, partition_key, queueing_task) triples in production + # (see _queue_pending_reservations); the task is irrelevant to this + # specific release path (only_current_task defaults False here). + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} async def detached_release(): return await limiter._pop_pending_concurrency_keys(model_call_details) released = await asyncio.create_task(detached_release()) - assert released == ("key1",) + assert released == (("key1", None),) # The parent's own view of the same dict must see the release too. assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [] @@ -4000,31 +4130,56 @@ async def test_release_in_a_forked_task_is_visible_to_the_parent_context(time_co @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"]} + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} async def detached_release_then_sibling_admits(): released = await limiter._pop_pending_concurrency_keys(model_call_details) # A sibling hop's admission, appending to the same shared dict, # interleaved right after this release's snapshot was taken. - model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD].append("key2") + model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD].append(("key2", None, None)) return released released = await asyncio.create_task(detached_release_then_sibling_admits()) - assert released == ("key1",) + assert released == (("key1", None),) # key2 must still be pending for its own hop's eventual release. - assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == ["key2"] + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("key2", None, None)] @pytest.mark.asyncio async def test_release_is_not_repeated_for_the_same_snapshot(time_controller): limiter = _make_limiter(time_controller) - model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: ["key1"]} + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} first = await limiter._pop_pending_concurrency_keys(model_call_details) second = await limiter._pop_pending_concurrency_keys(model_call_details) - assert first == ("key1",) + assert first == (("key1", None),) assert second == () +@pytest.mark.asyncio +async def test_release_only_current_task_leaves_a_concurrent_siblings_reservation_alone(time_controller): + """ + Veria AI finding: Router.abatch_completion's comma-separated multi-model + dispatch runs each model concurrently as its own asyncio.Task, but every + branch shares one litellm_logging_obj (the proxy attaches it to the + request before the comma-split), so a new hop's admission could see a + still-live sibling branch's own reservation sitting in the same + model_call_details and wrongly sweep it up as "stale". only_current_task + must leave a differently-tasked entry untouched. + """ + limiter = _make_limiter(time_controller) + + async def _reserve_as_a_separate_task() -> None: + pass # the task object itself is the fixture; body is irrelevant + + sibling_task = asyncio.create_task(_reserve_as_a_separate_task()) + await sibling_task + model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("sibling-key", None, sibling_task)]} + + released = await limiter._pop_pending_concurrency_keys(model_call_details, only_current_task=True) + assert released == () + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_task)] + + # --------------------------------------------------------------------------- # refund-on-rollback across differently-hash-tagged keys (Redis Cluster safety) # --------------------------------------------------------------------------- From 744733ac2d79186241d4d5bc38542b3627145086 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 18:38:42 -0400 Subject: [PATCH 25/44] chore(proxy): regenerate the stale lazy OpenAPI snapshot and dashboard schema.d.ts CI's "Check UI API Types Sync" job failed on d6d68f866e: the frozen litellm/proxy/_lazy_openapi_snapshot.json (used to serve /openapi.json for routes not imported in the current process) had drifted out of sync with the actual lazily-loaded routes/models already present in this branch's code (e.g. MCPCredentials.upstream_token_header, a2a_registration). Unrelated to this PR's own tag-rate-limiting changes -- regenerating via `python -m litellm.proxy._lazy_openapi_snapshot` then `npm run gen:api` (both commands CI's own failure message names) brings both frozen artifacts back in sync with the code already on this branch. --- litellm/proxy/_lazy_openapi_snapshot.json | 35 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +++ 2 files changed, 40 insertions(+) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 1963c7799a2..040d258f97a 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -10238,6 +10238,18 @@ "description": "AWS Bedrock runtime endpoint URL", "title": "Aws Bedrock Runtime Endpoint" }, + "aws_external_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "External ID required by the target role's trust policy on sts:AssumeRole", + "title": "Aws External Id" + }, "aws_profile_name": { "anyOf": [ { @@ -25237,6 +25249,9 @@ }, { "$ref": "#/components/schemas/ChatCompletionImageObject" + }, + { + "$ref": "#/components/schemas/ChatCompletionToolReferenceObject" } ] }, @@ -25324,6 +25339,26 @@ "title": "ChatCompletionToolParamFunctionChunk", "type": "object" }, + "ChatCompletionToolReferenceObject": { + "description": "Anthropic tool-search result block, carried through untouched so it survives a round trip.", + "properties": { + "tool_name": { + "title": "Tool Name", + "type": "string" + }, + "type": { + "const": "tool_reference", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "tool_name" + ], + "title": "ChatCompletionToolReferenceObject", + "type": "object" + }, "ChatCompletionUserMessage": { "properties": { "cache_control": { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9a968e53fde..511cf3c0013 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29466,6 +29466,11 @@ export interface components { * @description AWS Bedrock runtime endpoint URL */ aws_bedrock_runtime_endpoint?: string | null; + /** + * Aws External Id + * @description External ID required by the target role's trust policy on sts:AssumeRole + */ + aws_external_id?: string | null; /** * Aws Profile Name * @description AWS profile name for credential retrieval From c2d881894e803004db86b24186a2e7d21993e67d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 18:56:00 -0400 Subject: [PATCH 26/44] fix(proxy): stop a batch sibling's terminal event from releasing another branch's live concurrency slot Bugbot finding: async_log_success_event/async_log_failure_event still called _pop_pending_concurrency_keys with no filter, so the first finishing branch of an abatch_completion dispatch released every reservation on the shared model_call_details, including a still-live sibling branch's own slot. Reservations are now tagged with an admission-scoped token from a ContextVar rather than a raw asyncio.Task: create_task snapshots the current Context, so a task explicitly forked from within one hop's own admission (litellm's own logging dispatch explicitly propagates context) still reads back the same token, while abatch_completion's sibling branches, forked before any of them call admission, each mint their own distinct token on first use. Both the stale-hop cleanup and the terminal release hooks (success/failure) now require this token to match before releasing; the disconnect hook is left unfiltered, since its own scenario (mid-stream disconnect) cannot co-occur with abatch_completion's combined, non-streaming response. --- .../hooks/model_based_tag_rate_limits_hook.py | 134 +++++++---- .../test_model_based_tag_rate_limits_hook.py | 209 ++++++++++++++---- 2 files changed, 256 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 4c32f822b42..a1bba1af4d2 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -12,6 +12,7 @@ per-deployment dedup. """ import asyncio +import contextvars import json from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, replace @@ -558,15 +559,42 @@ _INDEX_TTL_SECONDS: Final = 5.0 # can't be forged or guessed. # # "shared across that request's own fallback hops" is not the same as -# "scoped to one asyncio Task": `Router.abatch_completion`'s comma-separated -# multi-model dispatch runs several branches concurrently, each its own -# Task, but hands every branch the identical `litellm_logging_obj` -- so -# each entry also carries the Task that queued it (see -# `_queue_pending_reservations`), letting `_release_stale_hop_reservations` -# tell a genuinely stale same-task hop apart from a still-live sibling -# branch's own reservation. +# "shared across every task that happens to touch model_call_details": +# `Router.abatch_completion`'s comma-separated multi-model dispatch runs +# several branches concurrently, each its own Task, but hands every branch +# the identical `litellm_logging_obj` -- so a still-live sibling branch's +# own reservation can sit in this same list. Each entry also carries an +# admission-scoped token (see `_current_admission_token`) so a release can +# tell a genuinely stale same-lineage hop apart from a still-live sibling +# branch's own reservation, without reintroducing the non-descendant-task +# blind spot a bare `ContextVar` has for the data itself (see above): the +# token is only ever compared for *identity*, never relied on to carry the +# reservation across a task boundary the way `model_call_details` does. _PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys" +# Identifies which admission call queued a given reservation, scoped by +# asyncio Context rather than by Task identity: `asyncio.create_task` +# snapshots the *current* Context into the new Task, so a task explicitly +# forked from within one hop's own admission (e.g. litellm's own +# `LoggingWorker.enqueue` dispatch, which explicitly propagates the calling +# context) still reads back the same token, while `abatch_completion`'s +# sibling branches -- forked *before* any of them ever called admission -- +# each start from an unset ContextVar and mint their own distinct token on +# first use, never matching each other's. +_ADMISSION_CONTEXT: Final[contextvars.ContextVar[object | None]] = contextvars.ContextVar( + "_model_based_tag_rate_limits_admission_context", default=None +) + + +def _current_admission_token() -> object: + token: Final = _ADMISSION_CONTEXT.get() + if token is not None: + return token + fresh_token: Final = object() + _ADMISSION_CONTEXT.set(fresh_token) + return fresh_token + + # Same `model_call_details`-stashing rationale as the field above, for a # different unit: "requests" is atomic and admitted once per hop (see # _ATOMIC_UNITS), same as concurrency, but a "requests" limit is meant to cap @@ -898,20 +926,20 @@ def _queue_pending_reservations( ) -> 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. + not a `litellm_call_id`. Silently a no-op without a real logging object + (defensive only; every real request has one): a queued concurrency + reservation still self-heals via `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, + just later. - Each entry is stamped with the queueing coroutine's own `asyncio.Task`: - `Router.abatch_completion`'s comma-separated multi-model dispatch runs - several `acompletion` calls concurrently as *separate* tasks that all - share one `model_call_details` (the proxy attaches one `litellm_logging_obj` - to the request before the comma-split, and every branch inherits that - same reference), so this field is no longer scoped to one logical - request's own serial fallback chain the way its docstring assumes. - `_release_stale_hop_reservations` uses the stamp to tell "an earlier hop - of *this* chain, safe to reclaim" apart from "a concurrent sibling + Each entry is stamped with `_current_admission_token()`: `Router.abatch_completion`'s + comma-separated multi-model dispatch runs several `acompletion` calls + concurrently as *separate* tasks that all share one `model_call_details` + (the proxy attaches one `litellm_logging_obj` to the request before the + comma-split, and every branch inherits that same reference), so this + field is no longer scoped to one logical request's own serial fallback + chain the way its docstring assumes. `_release_stale_hop_reservations` + and the terminal release hooks use the stamp to tell "this same + admission lineage, safe to reclaim" apart from "a concurrent sibling branch's own still-live reservation," which must never be touched here. """ logging_obj: Final = request_kwargs.get("litellm_logging_obj") @@ -922,9 +950,9 @@ def _queue_pending_reservations( 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 - current_task: Final = asyncio.current_task() + admission_token: Final = _current_admission_token() pending.extend( - (key, partition_key, current_task) for key, partition_key in reservations + (key, partition_key, admission_token) for key, partition_key in reservations ) # mutable-ok: see comment above @@ -1483,24 +1511,24 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] async def _release_stale_hop_reservations(self, request_kwargs: Mapping[str, object]) -> frozenset[str]: """ A concurrency reservation still queued when a *new* hop's admission - runs, *within the same asyncio Task*, can only belong to an earlier - hop of this same request's own fallback chain that already concluded - and failed: Router awaits one hop's entire attempt (call plus its own - failure handling) before starting the next, and a hop that instead - succeeded ends the request there via async_log_success_event, which - already pops everything -- so admission is never re-entered, in that - same task, while an earlier hop's reservation is still legitimately - in flight. + runs, *within the same admission lineage*, can only belong to an + earlier hop of this same request's own fallback chain that already + concluded and failed: Router awaits one hop's entire attempt (call + plus its own failure handling) before starting the next, and a hop + that instead succeeded ends the request there via + async_log_success_event, which already pops its own lineage's + entries -- so admission is never re-entered, in that same lineage, + while an earlier hop's reservation is still legitimately in flight. - The task check matters because `model_call_details` is not always + The lineage check matters because `model_call_details` is not always scoped to one such chain: `Router.abatch_completion`'s comma-separated multi-model dispatch runs several branches concurrently, each its own Task, but every branch shares the identical `litellm_logging_obj` (see `_queue_pending_reservations`'s own docstring) -- so a reservation queued by a still-running sibling branch can be sitting here too, and releasing it out from under that branch would let more calls through - a concurrency limit than it allows. Only entries this exact Task - queued are safe to treat as stale; anything else is left for its own + a concurrency limit than it allows. Only entries `_current_admission_token()` + stamped are safe to treat as stale; anything else is left for its own branch to release. LiteLLM only invokes a request's CustomLogger.async_log_failure_event @@ -1541,17 +1569,17 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] model_call_details: Final = getattr(logging_obj, "model_call_details", None) if not isinstance(model_call_details, dict): return frozenset() - release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_current_task=True) + release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_own_lineage=True) if release_keys: await self._release_keys(release_keys) pending_request_increments: Final = model_call_details.get(_PENDING_REQUEST_INCREMENTS_FIELD) if not isinstance(pending_request_increments, list): return frozenset() - current_task: Final = asyncio.current_task() - return frozenset(key for key, _partition_key, task in pending_request_increments if task is current_task) + admission_token: Final = _current_admission_token() + return frozenset(key for key, _partition_key, token in pending_request_increments if token is admission_token) async def _pop_pending_concurrency_keys( - self, kwargs: Mapping[str, object], *, only_current_task: bool = False + self, kwargs: Mapping[str, object], *, only_own_lineage: bool = False ) -> 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 @@ -1585,15 +1613,20 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # model_call_details can still be live and appending concurrently # (see the field's own docstring), so wiping the whole list here # would silently strand that branch's reservation instead of - # releasing it later. `only_current_task` additionally excludes any - # entry a *different*, still-running Task queued -- see + # releasing it later. `only_own_lineage` additionally excludes any + # entry a *different* admission lineage queued -- see # `_release_stale_hop_reservations`'s own docstring for why that - # distinction, not just presence, decides what's actually stale. + # distinction, not just presence, decides what's actually stale; + # the same distinction applies to a terminal release (success/failure) + # racing a still-live sibling branch's own reservation. pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) if not isinstance(pending, list) or not pending: return () - current_task: Final = asyncio.current_task() - snapshot: Final = tuple(entry for entry in pending if not only_current_task or entry[2] is current_task) + snapshot: Final = ( + tuple(entry for entry in pending if entry[2] is _current_admission_token()) + if only_own_lineage + else tuple(pending) + ) for entry in snapshot: try: pending.remove(entry) @@ -1610,6 +1643,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] Without this, the reservation would sit held until _CONCURRENCY_MIN_SAFETY_TTL_SECONDS expires, letting a caller who repeatedly opens and immediately drops streaming requests exhaust their own tag's concurrency limit for free. + + Deliberately not `only_own_lineage=True`: `Router.abatch_completion` + returns every branch's response together rather than a single stream, + so its concurrent-sibling-branch race this hook otherwise guards + against (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring) cannot + co-occur with a mid-stream disconnect here. """ logging_obj: Final = request_data.get("litellm_logging_obj") model_call_details: Final = getattr(logging_obj, "model_call_details", None) @@ -1685,13 +1724,18 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # 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) + # same request, and that slot must still be released. only_own_lineage + # keeps this from releasing a still-live sibling branch's own + # reservation when model_call_details is shared across an + # abatch_completion dispatch -- see _PENDING_CONCURRENCY_KEYS_FIELD's + # docstring. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs, only_own_lineage=True) 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) + # only_own_lineage: see async_log_failure_event's own comment above. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs, only_own_lineage=True) 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 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 index 60df49728f1..67a8c34bf29 100644 --- 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 @@ -24,6 +24,7 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( _build_group_limits, _build_limits_index, _ConfiguredLimit, + _current_admission_token, _extract_team_id, _inflight_key, _pending_reservations_cache_key, @@ -33,7 +34,14 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( BACKGROUND_TASKS as _BACKGROUND_TASKS, CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, ) -from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, RoutingGroup, TagRateLimitEntry, TagRateLimitScope +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + RoutingGroup, + TagRateLimitEntry, + TagRateLimitScope, +) class TimeController: @@ -230,7 +238,11 @@ async def test_filter_deployments_ignores_a_forged_empty_litellm_metadata_key(ti deployment = _deployment( "grp", "dep-1", - {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, ) router = litellm.Router(model_list=[deployment]) limiter.update_variables(llm_router=router) @@ -258,7 +270,11 @@ async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_presen deployment = _deployment( "grp", "dep-1", - {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, ) router = litellm.Router(model_list=[deployment]) limiter.update_variables(llm_router=router) @@ -288,7 +304,11 @@ async def test_filter_deployments_ignores_a_forged_populated_litellm_metadata_ke deployment = _deployment( "grp", "dep-1", - {"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}}, + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, ) router = litellm.Router(model_list=[deployment]) limiter.update_variables(llm_router=router) @@ -1730,7 +1750,15 @@ async def test_log_success_event_accounts_against_the_same_bucket_admission_chec 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 + "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 @@ -1772,9 +1800,11 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me model="my-group", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs ) assert admitted == healthy - admission_bucket_group = limiter._index.get(router).resolve_any( - "my-group", team_id=None, candidate_model_names=("backend-a", "backend-b") - )[0].resolved_group + admission_bucket_group = ( + limiter._index.get(router) + .resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b"))[0] + .resolved_group + ) routing_group = router.get_routing_group("my-group") assert routing_group is not None @@ -1799,7 +1829,15 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me now = time_controller.now().timestamp() admission_key = _expected_bucket_key( - "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group=admission_bucket_group, limit=500000 + "my-group", + "tokens", + "daily", + "end_user_id", + "u1", + 86400, + now, + resolved_group=admission_bucket_group, + limit=500000, ) drifted_key = _expected_bucket_key( "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-0", limit=500000 @@ -1881,7 +1919,13 @@ async def test_log_success_event_accounts_against_the_key_hash_admission_checked token_limits = { "token_limits": { "limits": [ - {"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400, "scope_by_key_hash": True} + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 500000, + "period_seconds": 86400, + "scope_by_key_hash": True, + } ] } } @@ -1933,7 +1977,9 @@ async def test_log_success_event_charges_the_window_admission_checked_not_a_late current when the response finishes. """ token_limits = { - "token_limits": {"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 500, "period_seconds": 60}]} + "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) @@ -1967,7 +2013,9 @@ async def test_log_success_event_charges_the_window_admission_checked_not_a_late ) assert ( float( - await limiter.internal_usage_cache.async_get_cache(key=admitted_window_bucket, litellm_parent_otel_span=None) + await limiter.internal_usage_cache.async_get_cache( + key=admitted_window_bucket, litellm_parent_otel_span=None + ) ) == 42.0 ) @@ -1993,7 +2041,11 @@ async def test_log_success_event_accounts_against_the_team_id_admission_checked( deployment = _deployment( "real-model-name", "dep-1", - {"token_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500, "period_seconds": 86400}]}}, + { + "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" @@ -2958,6 +3010,63 @@ async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_ assert len(rejections) == 1 +@pytest.mark.asyncio +async def test_concurrent_batch_siblings_terminal_event_does_not_release_a_live_siblings_slot(time_controller): + """ + Bugbot finding: async_log_success_event/async_log_failure_event still + released every pending reservation unconditionally, so the first + finishing abatch_completion branch freed a still-live sibling branch's + own concurrency slot too, letting a third, unrelated caller admit past + a limit that branch was still genuinely occupying. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=2) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + branch_two_admitted = asyncio.Event() + + async def _branch_two() -> None: + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + branch_two_admitted.set() + # Still "in flight" while branch one below finishes and releases. + await asyncio.sleep(0.05) + + task_two = asyncio.create_task(_branch_two()) + + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + await branch_two_admitted.wait() + 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) + + # Only branch one's own slot was freed; branch two's is still live, so + # exactly one more caller fits before the limit of 2 is hit again. + 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"]}}, + ) + await task_two + + def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[ @@ -2966,7 +3075,9 @@ def _request_limit_router(limit: int) -> "litellm.Router": "dep-1", { "request_limits": { - "limits": [{"name": "per_period", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300}] + "limits": [ + {"name": "per_period", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300} + ] } }, ) @@ -3419,7 +3530,9 @@ def _redis_limiter(time_controller: TimeController): 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 + return _PROXY_ModelBasedTagRateLimitsHook( + internal_usage_cache=dual_cache, time_provider=time_controller.now + ), redis_cache @pytest.mark.asyncio @@ -3545,7 +3658,11 @@ async def test_redis_backed_token_admission_sees_increments_the_in_memory_cache_ _deployment( "grp", "dep-1", - {"token_limits": {"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 100, "period_seconds": 60}]}}, + { + "token_limits": { + "limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 100, "period_seconds": 60}] + } + }, ) ] ) @@ -3626,7 +3743,9 @@ async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_co # would make a real-time before/after comparison too slow to assert # on deterministically) with refresh_ttl=True, matching how a # concurrency check is actually admitted. - admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + admitted, _ = await limiter._check_and_increment_one( + cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True + ) assert admitted ttl_after_first_admission = await redis_cache.redis_async_client.ttl(key) assert ttl_after_first_admission > 0 @@ -3636,7 +3755,9 @@ async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_co # A second admission on the same still-live key, most of the way # through the first admission's ttl, must push the ttl back out to # the full window again, not leave it counting down toward zero. - admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + admitted, _ = await limiter._check_and_increment_one( + cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True + ) assert admitted ttl_after_second_admission = await redis_cache.redis_async_client.ttl(key) assert ttl_after_second_admission >= 2 @@ -4112,9 +4233,9 @@ def test_concurrency_ttl_floor_does_not_shorten_a_longer_period_seconds(): @pytest.mark.asyncio async def test_release_in_a_forked_task_is_visible_to_the_parent_context(time_controller): limiter = _make_limiter(time_controller) - # Entries are (key, partition_key, queueing_task) triples in production - # (see _queue_pending_reservations); the task is irrelevant to this - # specific release path (only_current_task defaults False here). + # Entries are (key, partition_key, admission_token) triples in production + # (see _queue_pending_reservations); the token is irrelevant to this + # specific release path (only_own_lineage defaults False here). model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]} async def detached_release(): @@ -4156,28 +4277,26 @@ async def test_release_is_not_repeated_for_the_same_snapshot(time_controller): @pytest.mark.asyncio -async def test_release_only_current_task_leaves_a_concurrent_siblings_reservation_alone(time_controller): +async def test_release_only_own_lineage_leaves_a_concurrent_siblings_reservation_alone(time_controller): """ - Veria AI finding: Router.abatch_completion's comma-separated multi-model - dispatch runs each model concurrently as its own asyncio.Task, but every - branch shares one litellm_logging_obj (the proxy attaches it to the - request before the comma-split), so a new hop's admission could see a - still-live sibling branch's own reservation sitting in the same - model_call_details and wrongly sweep it up as "stale". only_current_task - must leave a differently-tasked entry untouched. + Bugbot/Veria AI finding: Router.abatch_completion's comma-separated + multi-model dispatch runs each model concurrently, each its own asyncio + Task, but every branch shares one litellm_logging_obj (the proxy attaches + it to the request before the comma-split), so a still-live sibling + branch's own reservation can sit in the same model_call_details. + only_own_lineage must release this context's own entry while leaving a + differently-lineaged (sibling branch's) entry untouched. """ limiter = _make_limiter(time_controller) + own_token = _current_admission_token() + sibling_token = object() + model_call_details: dict = { + _PENDING_CONCURRENCY_KEYS_FIELD: [("own-key", None, own_token), ("sibling-key", None, sibling_token)] + } - async def _reserve_as_a_separate_task() -> None: - pass # the task object itself is the fixture; body is irrelevant - - sibling_task = asyncio.create_task(_reserve_as_a_separate_task()) - await sibling_task - model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("sibling-key", None, sibling_task)]} - - released = await limiter._pop_pending_concurrency_keys(model_call_details, only_current_task=True) - assert released == () - assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_task)] + released = await limiter._pop_pending_concurrency_keys(model_call_details, only_own_lineage=True) + assert released == (("own-key", None),) + assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_token)] # --------------------------------------------------------------------------- @@ -4230,7 +4349,9 @@ async def test_cross_unit_refund_leaves_no_phantom_increment_in_memory(time_cont ) now = time_controller.now().timestamp() - request_key = _expected_bucket_key("grp", "requests", "per_minute", "end_user_id", "refund-check", 60, now, limit=10) + 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 @@ -4304,7 +4425,9 @@ async def test_exception_mid_batch_refunds_every_earlier_admission_before_propag raising_key = "{tag_rl:test:exception-refund:b}:requests" class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook): - async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool): + async def _check_and_increment_one( + self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool + ): if key == raising_key: raise RuntimeError("simulated transient redis failure") return await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) @@ -4342,7 +4465,9 @@ async def test_a_raising_keys_own_ambiguous_outcome_is_never_refunded(time_contr raising_key = "{tag_rl:test:ambiguous-no-refund:b}:requests" class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook): - async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool): + async def _check_and_increment_one( + self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool + ): if key == raising_key: # Simulate Redis committing the increment before the # response is lost: the write actually happens... From 66f2aff868ca5af6f365d0e8b8c07418a9ddce4a Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 20:33:52 -0400 Subject: [PATCH 27/44] fix(proxy): stop restricting terminal concurrency release to admission's own context, preserve team-alias resolved_group through routing-group fallback, merge instead of overwrite the pending-reservations mirror Veria AI / Cursor Bugbot findings on the prior admission-scoped ContextVar fix: a streaming response's success/failure event fires from a task the proxy forks independently of admission's own, so it never carries the same context token, and the earlier fix silently stopped releasing every completed stream's concurrency reservation until the safety TTL. async_log_success_event/async_log_failure_event go back to releasing unconditionally; only _release_stale_hop_reservations, which only ever runs synchronously within admission's own hop sequence, keeps the token check. resolve_any's routing-group fallback also always restamped resolved_group with the candidate model_name, discarding the team_public_model_name _build_limits_index already stamped there for a team-owned deployment -- splitting one team's bucket depending on whether a call reached it via its alias or a routing group. Now left untouched when already set. The pending-reservations cache mirror (the async_post_call_failure_hook fallback for when model_call_details itself is unavailable) is keyed by (call_id, key_hash), both identical across every branch of one abatch_completion dispatch. Mirroring now merges onto whatever's already cached instead of overwriting it, and releasing removes only the specific entries actually released instead of deleting the whole key, so one branch's own write or release can no longer erase a still-live sibling branch's own mirrored reservation before it's ever read. --- .../hooks/model_based_tag_rate_limits_hook.py | 165 ++++++++++++------ .../test_model_based_tag_rate_limits_hook.py | 151 ++++++++++------ 2 files changed, 207 insertions(+), 109 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index a1bba1af4d2..b425571d0c2 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -393,6 +393,13 @@ class _LimitsIndex: workers resolving the identical candidate set could otherwise pick different members as `resolved_group` and end up checking/accounting against different Redis keys for what's meant to be one shared bucket. + + A candidate's own entry can already carry `resolved_group` set to its + team's `team_public_model_name` (see `_build_limits_index`) -- that + stamp is left untouched rather than overwritten with `name` here, or + a team-owned deployment reachable through *both* its team alias and a + routing group would hash to two different buckets depending on which + path a caller happened to take. """ direct: Final = self.resolve(model, team_id) if direct: @@ -414,7 +421,9 @@ class _LimitsIndex: limit.deployment_scope, limit.team_scope, ) - deduped.setdefault(key, replace(limit, resolved_group=name)) # mutable-ok: see docstring above + deduped.setdefault( # mutable-ok: see docstring above + key, limit if limit.resolved_group is not None else replace(limit, resolved_group=name) + ) return tuple(deduped.values()) @@ -564,12 +573,18 @@ _INDEX_TTL_SECONDS: Final = 5.0 # several branches concurrently, each its own Task, but hands every branch # the identical `litellm_logging_obj` -- so a still-live sibling branch's # own reservation can sit in this same list. Each entry also carries an -# admission-scoped token (see `_current_admission_token`) so a release can -# tell a genuinely stale same-lineage hop apart from a still-live sibling -# branch's own reservation, without reintroducing the non-descendant-task -# blind spot a bare `ContextVar` has for the data itself (see above): the -# token is only ever compared for *identity*, never relied on to carry the -# reservation across a task boundary the way `model_call_details` does. +# admission-scoped token (see `_current_admission_token`), but that token is +# only ever trusted at the *next admission's own* stale-hop cleanup (see +# `_release_stale_hop_reservations`), never at the terminal release hooks +# below: a `ContextVar` has the identical non-descendant-task blind spot +# documented above for the reservation data itself, so a streaming +# response's own success event -- fired from a task the proxy forked +# independently of admission's -- would never see a matching token either, +# leaking every streaming request's reservation instead of releasing it. +# The terminal hooks release unconditionally, accepting the narrower risk +# that a still-live `abatch_completion` sibling's own reservation gets +# freed early, in exchange for actually releasing the far more common +# single-branch (including streaming) case. _PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys" # Identifies which admission call queued a given reservation, scoped by @@ -580,7 +595,10 @@ _PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pend # context) still reads back the same token, while `abatch_completion`'s # sibling branches -- forked *before* any of them ever called admission -- # each start from an unset ContextVar and mint their own distinct token on -# first use, never matching each other's. +# first use, never matching each other's. Only safe where the reader is +# guaranteed to run in a descendant of the writer's task -- see +# `_PENDING_CONCURRENCY_KEYS_FIELD`'s own docstring for where that doesn't +# hold. _ADMISSION_CONTEXT: Final[contextvars.ContextVar[object | None]] = contextvars.ContextVar( "_model_based_tag_rate_limits_admission_context", default=None ) @@ -1348,12 +1366,25 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] async def _mirror_pending_reservations( self, call_id: object, key_hash: str | None, reservations: Sequence[tuple[str, "_PartitionKey"]] ) -> None: + # Merged onto whatever's already mirrored, not overwritten: `call_id` + # and `key_hash` are identical across every branch of one + # abatch_completion dispatch (see _PENDING_CONCURRENCY_KEYS_FIELD's + # docstring), so an overwrite here would erase a still-live sibling + # branch's own mirrored reservation the moment this hop admits. + # Read-then-write is not atomic against a concurrent sibling doing + # the identical merge, so a rare interleaving can still lose one + # side's addition -- narrower than the guaranteed loss an overwrite + # caused, and self-heals via the same TTL either way. if not isinstance(call_id, str): return + cache_key: Final = _pending_reservations_cache_key(call_id, key_hash) try: + existing: Final = _decode_reservations( + await self.internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None) + ) await self.internal_usage_cache.async_set_cache( - key=_pending_reservations_cache_key(call_id, key_hash), - value=_encode_reservations(reservations), + key=cache_key, + value=_encode_reservations(existing + tuple(reservations)), ttl=_CONCURRENCY_MIN_SAFETY_TTL_SECONDS, litellm_parent_otel_span=None, ) @@ -1516,9 +1547,10 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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 its own lineage's - entries -- so admission is never re-entered, in that same lineage, - while an earlier hop's reservation is still legitimately in flight. + async_log_success_event, which already pops everything pending on + this same model_call_details -- so admission is never re-entered, in + that same lineage, while an earlier hop's reservation is still + legitimately in flight. The lineage check matters because `model_call_details` is not always scoped to one such chain: `Router.abatch_completion`'s comma-separated @@ -1581,33 +1613,6 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] async def _pop_pending_concurrency_keys( self, kwargs: Mapping[str, object], *, only_own_lineage: bool = False ) -> 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_authoritative_metadata_variable_name`'s own docstring. - litellm_params_raw: Final = kwargs.get("litellm_params") - litellm_params_for_metadata: Final = ( - litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs - ) - metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) - key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) - 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 entries, never a blanket # clear: a sibling branch sharing this same request's # model_call_details can still be live and appending concurrently @@ -1616,9 +1621,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # releasing it later. `only_own_lineage` additionally excludes any # entry a *different* admission lineage queued -- see # `_release_stale_hop_reservations`'s own docstring for why that - # distinction, not just presence, decides what's actually stale; - # the same distinction applies to a terminal release (success/failure) - # racing a still-live sibling branch's own reservation. + # distinction, not just presence, decides what's actually stale. pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) if not isinstance(pending, list) or not pending: return () @@ -1632,7 +1635,59 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] pending.remove(entry) except ValueError: pass - return tuple(entry[:2] for entry in snapshot) + released: Final = tuple(entry[:2] for entry in snapshot) + if released: + await self._discard_from_mirror(kwargs, released) + return released + + async def _discard_from_mirror( + self, kwargs: Mapping[str, object], released: Sequence[tuple[str, _PartitionKey]] + ) -> None: + # Removes only the entries this pop actually released, not a blanket + # delete of the whole (call_id, key_hash) mirror: that key is shared + # across every branch of one abatch_completion dispatch (see + # _PENDING_CONCURRENCY_KEYS_FIELD's docstring), so wiping it here + # would silently drop a still-live sibling branch's own entry before + # async_post_call_failure_hook ever gets to read it. + call_id: Final = kwargs.get("litellm_call_id") + if not isinstance(call_id, str): + return + # Not `get_metadata_variable_name_from_kwargs` (naive key-presence + # check): at this point `kwargs` is `model_call_details`, which + # carries `litellm_metadata` present-but-`None` alongside the real, + # populated `metadata` for a standard request -- see + # `_resolve_authoritative_metadata_variable_name`'s own docstring. + litellm_params_raw: Final = kwargs.get("litellm_params") + litellm_params_for_metadata: Final = litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) + key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) + cache_key: Final = _pending_reservations_cache_key(call_id, key_hash) + try: + mirrored: Final = list( # mutable-ok: local working copy, discarded after removing exactly `released`'s own occurrences below + _decode_reservations( + await self.internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None) + ) + ) + for entry in released: + try: + mirrored.remove(entry) + except ValueError: + pass + if mirrored: + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=_encode_reservations(tuple(mirrored)), + ttl=_CONCURRENCY_MIN_SAFETY_TTL_SECONDS, + litellm_parent_otel_span=None, + ) + else: + await self.internal_usage_cache.dual_cache.async_delete_cache(cache_key) + except Exception as e: # noqa: BLE001 - a failed mirror update must never block the real release above + verbose_proxy_logger.warning( + "model_based_tag_rate_limits_hook: failed to update mirrored reservations for call_id=%s: %s", + call_id, + e, + ) async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: """ @@ -1724,18 +1779,24 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # 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. only_own_lineage - # keeps this from releasing a still-live sibling branch's own - # reservation when model_call_details is shared across an - # abatch_completion dispatch -- see _PENDING_CONCURRENCY_KEYS_FIELD's - # docstring. - release_keys: Final = await self._pop_pending_concurrency_keys(kwargs, only_own_lineage=True) + # same request, and that slot must still be released. + # + # Not `only_own_lineage=True`: a streaming response is consumed (and + # this event fired) from a task the proxy forks independently of + # admission's own, so `_ADMISSION_CONTEXT` never reaches it either -- + # requiring a match here would leak every streaming request's + # reservation until `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS` instead of + # releasing it. `async_release_disconnect_state_hook` stays + # unfiltered for the identical reason; releasing everything here + # keeps a still-live `abatch_completion` sibling's own reservation + # exposed to the same risk that hook already accepts. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) if release_keys: await self._release_keys(release_keys) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: - # only_own_lineage: see async_log_failure_event's own comment above. - release_keys: Final = await self._pop_pending_concurrency_keys(kwargs, only_own_lineage=True) + # Not `only_own_lineage=True`: see async_log_failure_event's own comment above. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) if release_keys: release_task: Final = asyncio.create_task(self._release_keys(release_keys)) _BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring 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 index 67a8c34bf29..3e34715a606 100644 --- 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 @@ -2673,6 +2673,72 @@ async def test_post_call_failure_hook_cannot_release_a_different_keys_reservatio assert result == healthy +@pytest.mark.asyncio +async def test_mirror_write_does_not_erase_a_concurrent_siblings_own_reservation(time_controller): + """ + Cursor Bugbot finding: Router.abatch_completion's branches share one + litellm_call_id and one authenticated key, so they mirror their own + concurrency reservation under the identical (call_id, key_hash) cache + entry. Mirroring used to overwrite rather than merge, so branch two's own + admission silently erased branch one's already-mirrored reservation -- + branch one's own eventual chain-exhausting failure would then find + nothing to release there, permanently leaking its slot (until the safety + TTL) even though both branches genuinely finished. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=2) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + shared_call_id = "batch-call-id" + shared_tags = {"tags": ["end_user_id:shared-user"], "user_api_key": "batch-key-hash"} + branch_one_kwargs = {"metadata": shared_tags, "litellm_call_id": shared_call_id} + branch_two_kwargs = {"metadata": shared_tags, "litellm_call_id": shared_call_id} + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=branch_one_kwargs + ) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=branch_two_kwargs + ) + + # Both branches finish (order doesn't matter -- whichever hook reads the + # mirror first releases whatever is genuinely still there for either). + await limiter.async_post_call_failure_hook( + request_data={"litellm_call_id": shared_call_id}, + original_exception=Exception("branch one exhausted its fallbacks"), + user_api_key_dict=UserAPIKeyAuth(api_key="batch-key-hash"), + ) + await limiter.async_post_call_failure_hook( + request_data={"litellm_call_id": shared_call_id}, + original_exception=Exception("branch two exhausted its fallbacks"), + user_api_key_dict=UserAPIKeyAuth(api_key="batch-key-hash"), + ) + + # Neither slot leaked: two fresh admissions fit under the limit of 2, + # and a third does not. Under the erasure bug, branch two's write wiped + # branch one's mirrored entry before it was ever read, so only one of + # the two reservations was ever actually released. + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared-user"]}}, + ) + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared-user"]}}, + ) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:shared-user"]}}, + ) + + @pytest.mark.asyncio async def test_concurrency_slot_released_when_a_different_hook_rejects_the_request(time_controller): """ @@ -3010,63 +3076,6 @@ async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_ assert len(rejections) == 1 -@pytest.mark.asyncio -async def test_concurrent_batch_siblings_terminal_event_does_not_release_a_live_siblings_slot(time_controller): - """ - Bugbot finding: async_log_success_event/async_log_failure_event still - released every pending reservation unconditionally, so the first - finishing abatch_completion branch freed a still-live sibling branch's - own concurrency slot too, letting a third, unrelated caller admit past - a limit that branch was still genuinely occupying. - """ - limiter = _make_limiter(time_controller) - router = _concurrency_router(limit=2) - limiter.update_variables(llm_router=router) - healthy = router.model_list - request_kwargs, kwargs = _call_context(["end_user_id:u1"]) - branch_two_admitted = asyncio.Event() - - async def _branch_two() -> None: - await limiter.async_filter_deployments( - model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs - ) - branch_two_admitted.set() - # Still "in flight" while branch one below finishes and releases. - await asyncio.sleep(0.05) - - task_two = asyncio.create_task(_branch_two()) - - await limiter.async_filter_deployments( - model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs - ) - await branch_two_admitted.wait() - 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) - - # Only branch one's own slot was freed; branch two's is still live, so - # exactly one more caller fits before the limit of 2 is hit again. - 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"]}}, - ) - await task_two - - def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[ @@ -3854,6 +3863,34 @@ def test_build_limits_index_computes_identical_bucket_key_for_alias_and_internal assert key_via_internal_name == key_via_alias +def test_resolve_any_preserves_team_alias_resolved_group_through_routing_group_fallback(): + """ + Cursor Bugbot finding: resolve_any's routing-group fallback always + restamped resolved_group with the candidate's own model_name, discarding + the team_public_model_name _build_limits_index already stamped there -- + so a team-owned deployment reachable both via its team alias and via a + routing group split one team's usage across two buckets depending on + which path a caller took. + """ + deployment = _deployment( + "real-model-name", + "dep-1", + {"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}}, + ) + deployment["model_info"]["team_id"] = "team-1" + deployment["model_info"]["team_public_model_name"] = "team-alias-name" + index = _build_limits_index([deployment]) + + via_alias = index.resolve("team-alias-name", team_id="team-1")[0] + via_routing_group_fallback = index.resolve_any( + "my-group", team_id="team-1", candidate_model_names=["real-model-name"] + )[0] + + key_via_alias = _bucket_key("team-alias-name", via_alias, tag_value="u1", bucket_id=0) + key_via_fallback = _bucket_key("my-group", via_routing_group_fallback, tag_value="u1", bucket_id=0) + assert key_via_alias == key_via_fallback + + def test_build_limits_index_preserves_key_ttl_seconds_and_max_in_memory_cache_size(): """ Regression test: _configured_limit_for_signature used to reconstruct a From 4df02674760f4c5ab5684f757efe14316873d3b2 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:28:38 -0400 Subject: [PATCH 28/44] chore: retrigger CI (previous run's jobs were cancelled by infra/concurrency, not a real failure) From 332ba13f34bada3736b876942605b77e4b046b63 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:57:13 -0400 Subject: [PATCH 29/44] chore: retrigger CI (zizmor cancelled by infra/concurrency at the queue stage, not a real failure) From e3b6d5f7ba3b550387a45eb42b65d18e50514eaa Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:28:38 -0400 Subject: [PATCH 30/44] chore: retrigger CI (previous run's jobs were cancelled by infra/concurrency, not a real failure) From 6c792e659e17256905cece6b8d3464874f502cce Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:57:13 -0400 Subject: [PATCH 31/44] chore: retrigger CI (zizmor cancelled by infra/concurrency at the queue stage, not a real failure) From 243b5cd1ef5367101a7c197de357de41d4daa299 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 22:01:57 -0400 Subject: [PATCH 32/44] fix(proxy): release rate limit hook state on client disconnect A client disconnect throws GeneratorExit/CancelledError into the request path, so neither the success nor failure logging callback runs and a concurrency slot reserved at admission leaks until its own safety TTL. Gives every registered CustomLogger a chance to release such state via the new async_release_disconnect_state_hook, called from both the streaming and non-streaming cancel-on-disconnect paths. --- litellm/proxy/common_request_processing.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b582b164609..09d52b78da8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,7 +29,6 @@ from litellm.constants import ( NON_INFERENCE_CALL_TYPES, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, - STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -204,10 +203,6 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } -def _withheld_provider_output(response: object) -> bool: - return getattr(response, "has_buffered_provider_output", False) is True - - def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -3487,9 +3482,8 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. A keepalive ping carries no provider output, - # so it must not suppress that refund. - delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES + # False and refunds. + delivered_chunk = True yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): @@ -3503,7 +3497,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk and not _withheld_provider_output(response): + if not delivered_chunk: from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, ) From db0cbc9afbf8654a2edc276501cf8136a4133ec5 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:28:38 -0400 Subject: [PATCH 33/44] chore: retrigger CI (previous run's jobs were cancelled by infra/concurrency, not a real failure) From 1715bbc6904666bc802c745967ea3f03c9125910 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:57:13 -0400 Subject: [PATCH 34/44] chore: retrigger CI (zizmor cancelled by infra/concurrency at the queue stage, not a real failure) From e5ff0084b34780b8252b51f54d049dfd7038c8c5 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:28:38 -0400 Subject: [PATCH 35/44] chore: retrigger CI (previous run's jobs were cancelled by infra/concurrency, not a real failure) From 464a1638d9fbba36d1d78e7aefea6cd1cf4d0cd2 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:57:13 -0400 Subject: [PATCH 36/44] chore: retrigger CI (zizmor cancelled by infra/concurrency at the queue stage, not a real failure) From 41824eda82b477a4917255e6c2d1bd308b70add2 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 10:32:54 -0400 Subject: [PATCH 37/44] feat(rate-limiting): add global tag rate limiting hook Adds global_tag_rate_limits_hook, a model-independent async_pre_call_hook that enforces tag rate limits across an entire fallback chain, plus a cross_model_scope fallback guard so a chain-wide rejection isn't silently retried against an unlisted fallback model. --- litellm/__init__.py | 3 + litellm/litellm_core_utils/litellm_logging.py | 26 + litellm/proxy/common_request_processing.py | 5 +- .../hooks/global_tag_rate_limits_hook.py | 789 +++++++++++++ .../hooks/test_global_tag_rate_limits_hook.py | 1005 +++++++++++++++++ .../proxy/test_common_request_processing.py | 59 + 6 files changed, 1886 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/hooks/global_tag_rate_limits_hook.py create mode 100644 tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py diff --git a/litellm/__init__.py b/litellm/__init__.py index f9348f68f1b..b82c918230d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -120,6 +120,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "dynamic_rate_limiter", "dynamic_rate_limiter_v3", "model_based_tag_rate_limits_hook", + "global_tag_rate_limits_hook", "langsmith", "prometheus", "otel", @@ -393,6 +394,8 @@ default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None model_based_tag_rate_limits_max_in_memory_cache_size: Optional[int] = None +global_tag_rate_limits: Optional["TagRateLimits"] = None +global_tag_rate_limits_max_in_memory_cache_size: Optional[int] = None model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4f7a7510f96..4003b5f4d9a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4496,6 +4496,23 @@ def _init_custom_logger_compatible_class( model_based_tag_rate_limits_hook_obj.update_variables(llm_router=llm_router) _in_memory_loggers.append(model_based_tag_rate_limits_hook_obj) return model_based_tag_rate_limits_hook_obj + elif logging_integration == "global_tag_rate_limits_hook": + from litellm.proxy.hooks.global_tag_rate_limits_hook import ( + _PROXY_GlobalTagRateLimitsHook, # pyright: ignore[reportPrivateUsage] # resolved by name like every other opt-in callback here + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_GlobalTagRateLimitsHook): + return callback + + if internal_usage_cache is None: + raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}") + + global_tag_rate_limits_hook_obj: Final = _PROXY_GlobalTagRateLimitsHook( + internal_usage_cache=internal_usage_cache + ) + _in_memory_loggers.append(global_tag_rate_limits_hook_obj) + return global_tag_rate_limits_hook_obj elif logging_integration == "langtrace": if "LANGTRACE_API_KEY" not in os.environ: raise ValueError("LANGTRACE_API_KEY not found in environment variables") @@ -4945,6 +4962,15 @@ def get_custom_logger_compatible_class( if isinstance(callback, _PROXY_ModelBasedTagRateLimitsHook): return callback + elif logging_integration == "global_tag_rate_limits_hook": + from litellm.proxy.hooks.global_tag_rate_limits_hook import ( + _PROXY_GlobalTagRateLimitsHook, # pyright: ignore[reportPrivateUsage] # resolved by name like every other opt-in callback here + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_GlobalTagRateLimitsHook): + return callback + elif logging_integration == "langtrace": from litellm.integrations.opentelemetry import OpenTelemetry diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 09d52b78da8..0f54b50a16f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1998,7 +1998,10 @@ class ProxyBaseLLMRequestProcessing: ) except ProxyRateLimitError as original_exc: original_model: Final = self.data.get("model") - if not original_model or not llm_router or self.data.get("disable_fallbacks"): + cross_model_scope: Final = ( + isinstance(original_exc.detail, Mapping) and original_exc.detail.get("cross_model_scope") is True + ) + if not original_model or not llm_router or self.data.get("disable_fallbacks") or cross_model_scope: raise fallback_models: Final = self._resolve_fallback_models( diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py new file mode 100644 index 00000000000..ce8f59461d0 --- /dev/null +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -0,0 +1,789 @@ +""" +Tag-scoped token, request, dollar, and concurrency rate limits declared once, +globally, in `litellm_settings.global_tag_rate_limits` -- enforced once per +request in `async_pre_call_hook`, before Router does any routing, so a limit +applies regardless of which model or fallback chain the request ends up +hitting. + +This is the model-independent sibling of `model_based_tag_rate_limits_hook`, +which enforces the same `TagRateLimitEntry` shape but nested per-deployment +under `model_info.tag_rate_limits`, once per routing hop +(`async_filter_deployments`). A global entry has no deployment/routing-group +to reconcile -- there is exactly one config value, read once -- so this hook +reuses that sibling's free, already-hardened helper functions +(`entry_applies`, the Lua atomic check-and-increment scripts, cache +partitioning, bucket-key hashing primitives) from `tag_rate_limits_shared.py` +rather than duplicating them, but implements its own, much smaller +admission/accounting engine: no `_LimitsIndex`, no routing-group or +team-alias resolution, no per-deployment dedup signatures. + +Three independent entry-level knobs decide who a global entry applies to and +how its bucket is shared: + +- `apply_to_key_alias`: unset means every request, any key, any model. Set + to a list of virtual-key aliases, only those keys' requests count. +- `apply_to_models`: unset means every model. Set to a list of model names, + only requests whose caller-facing `model` field is in that list count -- + letting one entry rate-limit a whole fallback chain as a single unit by + naming every model in the chain. Each check is a fresh, independent + evaluation of `_entry_applies` against whatever `model` is current at that + moment, not a one-time decision that then sticks for the rest of the + request. Two concrete consequences follow from that: + (1) if the request's own model fails mid-flight and Router internally + retries a different model for the *same* admitted call, that retry is + never re-checked -- the original admission (against the originally + requested model) already stands, so an operator who needs the limit to + track whichever model actually ends up serving a request needs + `model_info.tag_rate_limits` instead; but + (2) if this hook's own admission *rejects* the request, + `common_request_processing.py` would otherwise catch that rejection and + retry the whole pre-call pipeline against + `litellm_settings.fallbacks`/`router_settings.fallbacks`, with + `data["model"]` mutated to the fallback target -- silently admitting the + request via a model outside `apply_to_models`, defeating the cap. A + rejection from an `apply_to_models`-scoped entry carries + `detail["cross_model_scope"] = True` for exactly this reason: + `_pre_call_with_fallbacks` checks that marker and re-raises immediately + instead of trying any fallback, so this bypass is closed regardless of + whether the fallback chain is also listed in `apply_to_models`. +- `scope_by_key_hash` (already exists on `TagRateLimitEntry`): whether the + keys an entry applies to share one bucket, or each gets its own. + +`async_pre_call_hook` runs before Router constructs `Logging`/`litellm_logging_obj` +for this request (see `common_request_processing.py`: `pre_call_hook` fires +well before `base_process_llm_request` builds the logging object), so unlike +`model_based_tag_rate_limits_hook` this hook cannot stash pending concurrency +reservations on `data["litellm_logging_obj"].model_call_details` -- that +object doesn't exist yet. Per-request state is instead kept on a +`ContextVar`-based stash, the same established pattern +`parallel_request_limiter_v3.py`'s v3 handler already uses for exactly this +problem, with one difference: the stash here is a dict keyed by +`litellm_call_id` rather than one shared mutable instance with an +overwritable "owner" field, so a nested LiteLLM call made inside the request +(e.g. a guardrail's own LLM judge call) -- which mints its own fresh call id +but inherits the same ContextVar-held ancestor context, not a separate one +-- gets its own isolated entry instead of overwriting the outer call's and +having its own success callback release the outer call's still-pending +reservations early. +""" + +import asyncio +from collections.abc import Callable, Mapping, Sequence +from contextvars import ContextVar +from dataclasses import dataclass, field +from datetime import datetime +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching model_based_tag_rate_limits_hook's identical import +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + ATOMIC_UNITS as _ATOMIC_UNITS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + BACKGROUND_TASKS as _BACKGROUND_TASKS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + LIMIT_UNITS as _LIMIT_UNITS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + TAG_RL_CHECK_AND_INCR_SCRIPT, + TAG_RL_DECR_FLOOR_ZERO_SCRIPT, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + UNIT_TO_GROUP_FIELD as _UNIT_TO_GROUP_FIELD, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + UNIT_TO_RATE_LIMIT_TYPE as _UNIT_TO_RATE_LIMIT_TYPE, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + LimitUnit as _LimitUnit, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + PartitionKey as _PartitionKey, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + PartitionOperations as _PartitionOperations, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + bucket_ttl_seconds as _bucket_ttl_seconds, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + entry_applies as _entry_applies, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + extract_identity as _extract_identity, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + extract_key_alias as _extract_key_alias, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + extract_key_hash as _extract_key_hash, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + fixed_length_identity as _fixed_length_identity, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + partition_key as _partition_key, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + policy_fingerprint as _policy_fingerprint, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + resolve_success_event_metadata_variable_name as _resolve_success_event_metadata_variable_name, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.router_strategy.tag_based_routing import ( + _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching model_based_tag_rate_limits_hook's identical import +) +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.router import TagRateLimitEntry, TagRateLimits +from litellm.types.utils import StandardLoggingPayload + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span: TypeAlias = _Span +else: + Span: TypeAlias = object + + +def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: + """ + Global-hook equivalent of `model_based_tag_rate_limits_hook._hash_tag`, + without a `model_group`/deployment-scope/team-scope dimension -- a global + entry has none of those. Namespaced under `tag_rl:global:` so it can never + collide with that sibling hook's own `tag_rl:{model_group}:...` keys even + if an operator names a deployment "global": every key also differs by + `unit`/`name`/`tag_id`/`_policy_fingerprint`, and the two hooks' entries + are never meant to share a bucket in the first place. + """ + key_suffix: Final = f":key:{key_hash}" if key_hash is not None else "" + policy_suffix: Final = f":policy:{_policy_fingerprint(entry)}" + return f"tag_rl:global:{unit}:{entry.name}:{entry.tag_id}:{_fixed_length_identity(tag_value)}{key_suffix}{policy_suffix}" + + +def _bucket_key( + entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, bucket_id: int, key_hash: str | None +) -> str: + return f"{{{_hash_tag(entry, unit, tag_value, key_hash)}}}:{bucket_id}" + + +def _inflight_key(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: + return f"{{{_hash_tag(entry, unit, tag_value, key_hash)}}}:inflight" + + +@dataclass(frozen=True, slots=True) +class _ClassifiedGlobalCheck: + unit: _LimitUnit + entry: TagRateLimitEntry + tag_value: str + key: str + is_atomic: bool + + +@dataclass(frozen=True, slots=True) +class _CachePartition: + internal_usage_cache: InternalUsageCache + v3: _PROXY_MaxParallelRequestsHandler_v3 + + +@dataclass(slots=True) +class _GlobalTagRateLimitStash: + """Per-call bookkeeping `async_pre_call_hook` hands to that same call's + success/failure/disconnect callbacks -- see module docstring for why + this lives on a `ContextVar`, not `model_call_details`. + + Keyed by `litellm_call_id` in the dict below rather than one shared + mutable instance with an overwritable "owner" field: a nested LiteLLM + call made inside the request (an LLM-judge guardrail, a silent + experiment) that mints its own fresh call id runs inside the *same* + inherited context, not a separate one, so a single shared instance's + owner field would get reassigned to the nested call and its own + success callback would then release the outer call's still-pending + reservations early -- letting extra same-tag requests through while + the outer request is still genuinely in flight. Keying by call id + isolates each call's own reservations regardless of nesting. + """ + + admission_time: float | None = None + # The caller-facing `model` admission read from `data.get("model")`, so + # async_log_success_event's tokens/dollars accounting gates + # apply_to_models against the same, originally-requested model admission + # decided on -- not whatever model a later fallback actually served. + model: str | None = None + pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue + # "requests" keys already charged for this call_id -- veria-ai finding: + # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reruns the whole + # pre-call pipeline (this hook included) once per fallback model on ANY + # ProxyRateLimitError, not only one this hook itself raised, but reuses + # the same litellm_call_id (self.data is mutated in place, only `model` + # changes) across every attempt -- so this stash is the SAME object each + # time. A "requests" check matching an already-charged key here renews + # at zero net cost instead of charging a second unit for the same + # logical request; see async_pre_call_hook's own comment for how. + charged_request_keys: list[str] = field(default_factory=list) # mutable-ok: see comment above + # The server-authenticated key_hash (UserAPIKeyAuth.api_key) of whichever + # call first claimed this stash. litellm_call_id is caller-controlled via + # the x-litellm-call-id header (the exact forgery vector + # model_based_tag_rate_limits_hook's own pending-reservations mirror was + # hardened against earlier), so two unrelated requests sharing a + # caller-chosen id must not be allowed to "renew" each other's charge -- + # only a later admission carrying this same, authenticated key_hash may. + owner_key_hash: str | None = None + + +# Sentinel key for a call with no litellm_call_id at all (claim and lookup +# both fall back to this same key, so behavior for that degenerate case is +# unchanged: everything without a call id still shares one bucket). +_NO_CALL_ID: Final = "" + +_StashByCallId: TypeAlias = dict[ + str, _GlobalTagRateLimitStash +] # mutable-ok: per-call-id entries added over a request's lifetime, see class docstring + +_request_stash: Final[ContextVar[_StashByCallId | None]] = ContextVar( + "global_tag_rate_limits_request_stash", default=None +) + + +def _claim_stash_for_data(data: Mapping[str, object]) -> _GlobalTagRateLimitStash: + by_call_id: _StashByCallId | None = _request_stash.get() # rebind-ok: lazily initialized below if never set + if by_call_id is None: + by_call_id = {} # rebind-ok: see above # mutable-ok: see _StashByCallId + _request_stash.set(by_call_id) + owner_call_id: Final = data.get("litellm_call_id") + key: Final = owner_call_id if isinstance(owner_call_id, str) else _NO_CALL_ID + stash = by_call_id.get(key) # rebind-ok: reassigned just below when newly created + if stash is None: + stash = _GlobalTagRateLimitStash() # rebind-ok: see above + by_call_id[key] = stash # mutable-ok: see class docstring + return stash + + +def _stash_for_call(litellm_call_id: str | None) -> _GlobalTagRateLimitStash | None: + by_call_id: Final = _request_stash.get() + if by_call_id is None: + return None + key: Final = litellm_call_id if litellm_call_id is not None else _NO_CALL_ID + return by_call_id.get(key) + + +def _call_id_from_kwargs(kwargs: Mapping[str, object]) -> str | None: + call_id: Final = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def _resolve_max_in_memory_cache_size() -> int | None: + """Same shape as `model_based_tag_rate_limits_hook`'s own function, reading + this hook's own `litellm_settings` knob instead.""" + configured: Final = litellm.global_tag_rate_limits_max_in_memory_cache_size + if isinstance(configured, int) and not isinstance(configured, bool) and configured > 0: + return configured + if configured is not None: + verbose_proxy_logger.warning( + "global_tag_rate_limits_hook: global_tag_rate_limits_max_in_memory_cache_size=%r is not a positive " + "integer; falling back to the default in-memory cache size.", + configured, + ) + return None + + +class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # only referenced via the deferred import in litellm_logging.py's callback resolver; basedpyright doesn't trace that usage + CustomLogger +): + def __init__( + self, + internal_usage_cache: DualCache, + time_provider: Callable[[], datetime] | None = None, + ) -> None: + self._redis_cache: Final = internal_usage_cache.redis_cache + self._time_provider = time_provider or datetime.now + self._partitions: dict[_PartitionKey, _CachePartition] = {} # mutable-ok: lazily memoized; see _partition_for + self._partitions_lock = asyncio.Lock() + default_partition: Final = self._build_partition(_resolve_max_in_memory_cache_size()) + self._partitions[None] = default_partition + self.internal_usage_cache = default_partition.internal_usage_cache + self._lock = asyncio.Lock() + redis_cache: Final = self._redis_cache + self._check_and_incr_script = ( + redis_cache.async_register_script(TAG_RL_CHECK_AND_INCR_SCRIPT) if redis_cache is not None else None + ) + self._decr_floor_zero_script = ( + redis_cache.async_register_script(TAG_RL_DECR_FLOOR_ZERO_SCRIPT) if redis_cache is not None else None + ) + self._config_cache_key: object | None = None + self._config: TagRateLimits | None = None + + def _refresh_config(self) -> TagRateLimits | None: + """Re-validates `litellm.global_tag_rate_limits` whenever the object + identity changes (a config reload replaces it wholesale via + `setattr(litellm, key, value)`), so a hot-reloaded config takes effect + on the very next request with no staleness window and no TTL to tune.""" + raw: Final = getattr(litellm, "global_tag_rate_limits", None) + if raw is not self._config_cache_key: + self._config = TagRateLimits.model_validate(raw) if raw else None + self._config_cache_key = raw + return self._config + + def _build_partition(self, cache_size_override: int | None) -> _CachePartition: + dual_cache: Final = DualCache( + in_memory_cache=InMemoryCache(max_size_in_memory=cache_size_override), + redis_cache=self._redis_cache, + ) + cache: Final = InternalUsageCache(dual_cache=dual_cache) + return _CachePartition( + internal_usage_cache=cache, + v3=_PROXY_MaxParallelRequestsHandler_v3(cache, time_provider=self._time_provider), + ) + + async def _partition_for(self, partition_key: _PartitionKey) -> _CachePartition: + existing: Final = self._partitions.get(partition_key) + if existing is not None: + return existing + async with self._partitions_lock: + existing_after_lock: Final = self._partitions.get(partition_key) + if existing_after_lock is not None: + return existing_after_lock + cache_size_override: Final = partition_key[-1] if partition_key is not None else None + built: Final = self._build_partition(cache_size_override) + self._partitions[partition_key] = ( + built # mutable-ok: lazily memoized per distinct partition key, guarded by _partitions_lock above + ) + return built + + async def _check_and_increment_one( + self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int + ) -> tuple[bool, float]: + if self._check_and_incr_script is not None: + raw: Final = await self._check_and_incr_script(keys=(key,), args=(limit, increment, ttl)) + return bool(raw[0]), float(raw[1]) + async with self._lock: + current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) + current: Final = float(current_value) if current_value is not None else 0.0 + if current + increment > limit: + return False, current + new_value: Final = current + increment + await cache.async_set_cache(key=key, value=new_value, ttl=ttl, litellm_parent_otel_span=None) + return True, new_value + + async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None: + if self._decr_floor_zero_script is not None: + await self._decr_floor_zero_script(keys=(key,), args=(delta,)) + return + async with self._lock: + current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) + current: Final = float(current_value) if current_value is not None else 0.0 + await cache.async_set_cache(key=key, value=max(0.0, current + delta), litellm_parent_otel_span=None) + + async def _atomic_check_and_increment( + self, + checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], + ) -> tuple[int | None, tuple[float, ...]]: + """All-or-nothing atomic admission across `checks` -- see + `model_based_tag_rate_limits_hook._PROXY_ModelBasedTagRateLimitsHook._atomic_check_and_increment`'s + own docstring for the full rationale (refund-on-rollback, why a + raising key's own outcome is never refunded); identical logic, + duplicated rather than shared since it lives as instance methods + rather than free functions.""" + if not checks: + return None, () + admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection + for index, (cache, key, limit, increment, ttl) in enumerate(checks): + admitted = False + try: + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl) + finally: + if not admitted: + await self._refund_admitted(checks, up_to_index=index) + if admitted: + admitted_values.append(value) # mutable-ok: see accumulator comment above + continue + return index, (value,) + return None, tuple(admitted_values) + + async def _refund_admitted( + self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], up_to_index: int + ) -> None: + for refund_index in range(up_to_index): + refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index] + try: + await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment) + except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest + verbose_proxy_logger.warning( + "global_tag_rate_limits_hook: failed to refund %s on rollback: %s", refund_key, e + ) + + async def _release_keys(self, reservations: Sequence[tuple[str, _PartitionKey]]) -> None: + for key, partition_key in reservations: + try: + partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration + await self._decrement_floor_zero(partition.internal_usage_cache, key, -1.0) + except Exception as e: # noqa: BLE001 - releasing a slot must never raise into the caller's request path + verbose_proxy_logger.warning( + "global_tag_rate_limits_hook: failed to release concurrency slot %s: %s", key, e + ) + + @staticmethod + def _ttl_for(unit: _LimitUnit, entry: TagRateLimitEntry) -> int: + if unit == "concurrency": + requested_ttl: Final = entry.key_ttl_seconds if entry.key_ttl_seconds is not None else entry.period_seconds + return max(requested_ttl, _CONCURRENCY_MIN_SAFETY_TTL_SECONDS) + return _bucket_ttl_seconds(entry) + + def _classify( + self, + config: TagRateLimits, + tags: Sequence[str], + key_alias: str | None, + key_hash: str | None, + now: float, + model: str | None, + ) -> tuple[_ClassifiedGlobalCheck, ...]: + classified: Final = [] # mutable-ok: sequential accumulator, immediately frozen into a tuple below + for unit in _LIMIT_UNITS: + group = getattr(config, _UNIT_TO_GROUP_FIELD[unit]) + if group is None: + continue + for entry in group.limits: + tag_value = _extract_identity(tags, entry.tag_id) + if tag_value is None: + continue + if not _entry_applies(entry, tags, key_alias, model): + continue + effective_key_hash = key_hash if entry.scope_by_key_hash else None + if unit == "concurrency": + key = _inflight_key(entry, unit, tag_value, key_hash=effective_key_hash) + classified.append( + _ClassifiedGlobalCheck(unit, entry, tag_value, key, is_atomic=True) + ) # mutable-ok: see comment above + continue + bucket_id = int(now) // entry.period_seconds + key = _bucket_key(entry, unit, tag_value, bucket_id, key_hash=effective_key_hash) + classified.append( # mutable-ok: see comment above + _ClassifiedGlobalCheck(unit, entry, tag_value, key, is_atomic=unit in _ATOMIC_UNITS) + ) + return tuple(classified) + + async def _read_only_values( + self, read_only_checks: Sequence[_ClassifiedGlobalCheck], parent_otel_span: Span | None + ) -> tuple[float | None, ...]: + if not read_only_checks: + return () + indices_by_partition: Final[dict[_PartitionKey, list[int]]] = {} # mutable-ok: grouped, reassembled below + for index, check in enumerate(read_only_checks): + partition_key = _partition_key(check.entry) + indices = indices_by_partition.setdefault(partition_key, []) # mutable-ok: see above + indices.append(index) # mutable-ok: see comment above + values_by_index: Final[dict[int, float | None]] = {} # mutable-ok: see comment above + for partition_key, indices in indices_by_partition.items(): + partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration + keys = [read_only_checks[i].key for i in indices] # mutable-ok: async_batch_get_cache needs a real list + redis_cache = partition.internal_usage_cache.dual_cache.redis_cache + if redis_cache is not None: + redis_values: Mapping[str, object] = await redis_cache.async_batch_get_cache( + key_list=keys, parent_otel_span=parent_otel_span + ) + resolved = [redis_values.get(key) for key in keys] # mutable-ok: needs a real list + else: + current_values = await partition.internal_usage_cache.async_batch_get_cache( + keys=keys, parent_otel_span=parent_otel_span, local_only=True + ) + missing = [None] * len(keys) # mutable-ok: async_batch_get_cache requires a real list; see above + resolved = current_values if current_values is not None else missing + for i, value in zip(indices, resolved): + values_by_index[i] = value # mutable-ok: see comment above + return tuple(values_by_index[i] for i in range(len(read_only_checks))) + + def _raise_if_over_limit( + self, + read_only_checks: Sequence[_ClassifiedGlobalCheck], + current_values: Sequence[float | None], + model: str | None, + ) -> None: + for check, current_value in zip(read_only_checks, current_values): + current = float(current_value) if current_value is not None else 0.0 + if current < check.entry.limit: + continue + self._raise_over_limit(check.unit, check.entry, check.tag_value, model, current=current) + + def _raise_over_limit( + self, unit: _LimitUnit, entry: TagRateLimitEntry, tag_value: str, model: str | None, current: float + ) -> None: + verbose_proxy_logger.debug( + "global_tag_rate_limits_hook: OVER_LIMIT unit=%s name=%s tag_id=%s tag_value=%s current=%s limit=%s", + unit, + entry.name, + entry.tag_id, + tag_value, + current, + entry.limit, + ) + raise ProxyRateLimitError( + detail={ # mutable-ok: async_log_failure_event and generic proxy exception rendering branch on isinstance(exc.detail, dict) + "error": "tag_rate_limit_exceeded", + "type": unit, + "tag_id": entry.tag_id, + "tag_value": tag_value, + "limit_name": entry.name, + "limit": entry.limit, + "period_seconds": entry.period_seconds, + # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reads this: + # an apply_to_models entry caps an entire named chain as one unit, so + # retrying against a fallback model outside that list would silently + # defeat the very policy that just rejected this request. + **({"cross_model_scope": True} if entry.apply_to_models is not None else {}), # mutable-ok: see above + }, + headers={"retry-after": str(entry.period_seconds)}, # mutable-ok: same as detail + rate_limit_type=_UNIT_TO_RATE_LIMIT_TYPE[unit], + model=model, + llm_provider="litellm_proxy", + ) + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, # mutable-ok: must match CustomLogger.async_pre_call_hook's own base signature exactly + call_type: str, + ) -> dict: # mutable-ok: must match CustomLogger.async_pre_call_hook's own base signature exactly + config: Final = self._refresh_config() + if config is None: + return data + + # async_pre_call_hook fires once per request in the common case, but + # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks can re-run + # this same pipeline once per fallback model on any ProxyRateLimitError + # (not only one this hook raised) -- see charged_request_keys' own + # docstring for how a repeat run for the same call_id renews rather + # than re-charges both "requests" and "concurrency" checks below. + stash: Final = _claim_stash_for_data(data) + + metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(data) + tags: Final = _get_tags_from_request_kwargs(data, metadata_variable_name=metadata_variable_name) + key_alias: Final = user_api_key_dict.key_alias + key_hash: Final = user_api_key_dict.api_key + model: Final = data.get("model") if isinstance(data.get("model"), str) else None + + # Only a repeat admission carrying the SAME authenticated key_hash as + # whichever call first claimed this stash may renew its charges -- + # see owner_key_hash's own docstring for why a bare call_id match is + # not enough. First admission for this stash claims ownership here. + if stash.owner_key_hash is None: + stash.owner_key_hash = key_hash + renewal_allowed: Final = stash.owner_key_hash == key_hash + + now: Final = self._time_provider().timestamp() + stash.admission_time = now + stash.model = model + classified: Final = self._classify(config, tags, key_alias, key_hash, now, model) + if not classified: + return data + + read_only_checks: Final = tuple(c for c in classified if not c.is_atomic) + atomic_checks: Final = tuple(c for c in classified if c.is_atomic) + + current_values: Final = await self._read_only_values(read_only_checks, parent_otel_span=None) + self._raise_if_over_limit(read_only_checks, current_values, model) + + if atomic_checks: + atomic_partitions_list: Final = [] # mutable-ok: sequential async lookups, one per atomic_checks entry + for check in atomic_checks: + atomic_partitions_list.append( + await self._partition_for(_partition_key(check.entry)) + ) # mutable-ok: see comment above + atomic_partitions: Final = tuple(atomic_partitions_list) + already_reserved_concurrency_keys: Final = frozenset( + key for key, _partition_key in stash.pending_concurrency_keys + ) + failing_index, values = await self._atomic_check_and_increment( + tuple( + ( + partition.internal_usage_cache, + check.key, + check.entry.limit, + # A key already charged/reserved for this call_id (an + # earlier _pre_call_with_fallbacks attempt for the + # same logical request) renews at zero net cost + # instead of charging or reserving a second unit -- + # folded into this same all-or-nothing batch so a + # rollback here (some other check in the batch + # rejecting) refunds that zero-cost renewal as a + # genuine no-op, same reasoning as + # model_based_tag_rate_limits_hook's identical fix + # for its own per-hop retries. + 0.0 + if renewal_allowed + and ( + (check.unit == "requests" and check.key in stash.charged_request_keys) + or (check.unit == "concurrency" and check.key in already_reserved_concurrency_keys) + ) + else 1.0, + self._ttl_for(check.unit, check.entry), + ) + for partition, check in zip(atomic_partitions, atomic_checks) + ) + ) + if failing_index is not None: + failing_check: Final = atomic_checks[failing_index] + self._raise_over_limit( + failing_check.unit, failing_check.entry, failing_check.tag_value, model, current=values[0] + ) + + # Only genuinely new reservations, never a key already in + # already_reserved_concurrency_keys: that key's own check just + # renewed at zero net cost above, so re-adding it here would + # make release (which decrements once per queued entry) decrement + # twice for a counter that was only ever incremented once. + concurrency_reservations: Final = tuple( + (check.key, _partition_key(check.entry)) + for check in atomic_checks + if check.unit == "concurrency" and check.key not in already_reserved_concurrency_keys + ) + if concurrency_reservations: + stash.pending_concurrency_keys.extend(concurrency_reservations) # mutable-ok: see field's own docstring + + # Only recorded when renewal_allowed: an admission that didn't + # own this stash (a call_id collision from a different key_hash) + # must not contaminate the rightful owner's own renewal + # tracking, or a later, genuine fallback retry from the owner + # could wrongly treat the impostor's charge as its own and + # renew for free. + request_keys: Final = ( + tuple( + check.key + for check in atomic_checks + if check.unit == "requests" and check.key not in stash.charged_request_keys + ) + if renewal_allowed + else () + ) + if request_keys: + stash.charged_request_keys.extend(request_keys) # mutable-ok: see field's own docstring + + return data + + async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + stash: Final = _stash_for_call(_call_id_from_kwargs(request_data)) + if stash is None or not stash.pending_concurrency_keys: + return + release_keys: Final = tuple(stash.pending_concurrency_keys) + stash.pending_concurrency_keys.clear() + await self._release_keys(release_keys) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + # No special-case skip for this hook's own tag_rate_limit_exceeded + # rejection: that rejection never reaches the point where a + # concurrency reservation is queued (see async_pre_call_hook), so + # stash.pending_concurrency_keys is already empty in that case and + # the check below naturally no-ops. Skipping release based on the + # exception's error marker alone would be wrong here, since + # model_based_tag_rate_limits_hook raises the identical marker -- + # that rejection can land after this hook already reserved a slot + # for this same request, and that slot must still be released. + stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) + if stash is None or not stash.pending_concurrency_keys: + return + release_keys: Final = tuple(stash.pending_concurrency_keys) + stash.pending_concurrency_keys.clear() + await self._release_keys(release_keys) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) + if stash is not None and stash.pending_concurrency_keys: + release_keys: Final = tuple(stash.pending_concurrency_keys) + stash.pending_concurrency_keys.clear() + release_task: Final = asyncio.create_task(self._release_keys(release_keys)) + _BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + release_task.add_done_callback(_BACKGROUND_TASKS.discard) + + config: Final = self._refresh_config() + if config is None: + return + + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") + if standard_logging_object is None: + return + + # kwargs here is Logging.model_call_details, not the router's flat + # request kwargs admission sees: metadata/litellm_metadata are never + # top-level here, only nested under kwargs["litellm_params"] (see + # Logging.update_environment_variables). + litellm_params_for_metadata: Final = kwargs.get("litellm_params") or kwargs + metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata) + key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) + key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name) + + tags: Final = _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name) + if not tags: + return + + now: Final = ( + stash.admission_time + if stash is not None and stash.admission_time is not None + else self._time_provider().timestamp() + ) + model: Final = stash.model if stash is not None else None + increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType( + { + "tokens": float(standard_logging_object.get("total_tokens") or 0), + "dollars": float(standard_logging_object.get("response_cost") or 0), + } + ) + + operation_by_entry: Final = [] # mutable-ok: sequential accumulator over config groups, immediately used below + for unit in ("tokens", "dollars"): + group = getattr(config, _UNIT_TO_GROUP_FIELD[unit]) + if group is None: + continue + for entry in group.limits: + tag_value = _extract_identity(tags, entry.tag_id) + if tag_value is None: + continue + if not _entry_applies(entry, tags, key_alias, model): + continue + increment_value = increment_by_unit[unit] + if increment_value == 0: + continue + bucket_id = int(now) // entry.period_seconds + key_hash_for_entry = key_hash if entry.scope_by_key_hash else None + key = _bucket_key(entry, unit, tag_value, bucket_id, key_hash=key_hash_for_entry) + operation_by_entry.append( # mutable-ok: see comment above + ( + entry, + RedisPipelineIncrementOperation( + key=key, increment_value=increment_value, ttl=_bucket_ttl_seconds(entry) + ), + ) + ) + + if not operation_by_entry: + return + + operations_by_partition: Final[_PartitionOperations] = {} # mutable-ok: grouped by cache partition below + for entry, operation in operation_by_entry: + partition_key = _partition_key(entry) + operations = operations_by_partition.setdefault(partition_key, []) # mutable-ok: see above + operations.append(operation) # mutable-ok: see comment above + + for partition_key, group_operations in operations_by_partition.items(): + partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration + accounting_task = asyncio.create_task( # not Final: rebound each loop iteration + partition.v3.async_increment_tokens_with_ttl_preservation( + pipeline_operations=tuple(group_operations), parent_otel_span=None + ) + ) + _BACKGROUND_TASKS.add(accounting_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + accounting_task.add_done_callback(_BACKGROUND_TASKS.discard) diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py new file mode 100644 index 00000000000..428037986e5 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -0,0 +1,1005 @@ +""" +Unit tests for the global-scope, model-independent tag rate limiter. +""" + +import asyncio +from datetime import datetime, timedelta + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.global_tag_rate_limits_hook import ( + _PROXY_GlobalTagRateLimitsHook, +) + + +class TimeController: + def __init__(self): + self._current = datetime(2026, 1, 1, 0, 0, 0) + + def now(self) -> datetime: + return self._current + + def advance(self, seconds: float) -> None: + self._current += timedelta(seconds=seconds) + + +@pytest.fixture +def time_controller(): + return TimeController() + + +def _make_hook(time_controller: TimeController) -> _PROXY_GlobalTagRateLimitsHook: + return _PROXY_GlobalTagRateLimitsHook( + internal_usage_cache=DualCache(), + time_provider=time_controller.now, + ) + + +def _key(alias: str | None = None, api_key: str = "hash") -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key=api_key, key_alias=alias) + + +def _data(tags: list[str], call_id: str = "call-1") -> dict: + return {"metadata": {"tags": tags}, "litellm_call_id": call_id} + + +# --------------------------------------------------------------------------- +# No-op when unconfigured +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_no_op_when_no_config_set(time_controller, monkeypatch): + monkeypatch.setattr(litellm, "global_tag_rate_limits", None) + hook = _make_hook(time_controller) + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + assert result == _data(["end_user_id:u1"]) + + +@pytest.mark.asyncio +async def test_malformed_config_raises_at_first_use(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, "global_tag_rate_limits", {"dollar_limits": {"limits": [{"name": "bad", "limit": "not-a-number"}]}} + ) + hook = _make_hook(time_controller) + with pytest.raises(ValidationError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + + +# --------------------------------------------------------------------------- +# Global scope: applies to every key by default +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_request_limit_shared_across_keys_by_default(time_controller, monkeypatch): + """No apply_to_key_alias -> the entry is one shared bucket regardless of + which key made the request.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a"), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + # A different key, identical tag value, and a distinct call_id (a + # genuinely separate logical request, not a fallback retry of the same + # one) -- must be rejected too, proving the bucket is genuinely shared, + # not per-key by default. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-b"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_request_limit_is_independent_of_model(time_controller, monkeypatch): + """The hook never reads `data["model"]` for identity -- two different + "models" (irrelevant to this hook) must still share the same bucket.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + data_model_a = {**_data(["end_user_id:u1"]), "model": "gpt-4o"} + # A distinct call_id: this is a separate logical request, not the same + # one retrying against a different model via _pre_call_with_fallbacks. + data_model_b = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "claude-3"} + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=data_model_a, call_type="completion" + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=data_model_b, call_type="completion" + ) + + +# --------------------------------------------------------------------------- +# apply_to_key_alias -- narrows which keys an entry applies to +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_ignores_non_matching_keys(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + for _ in range(3): + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="other-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_enforces_for_the_listed_key(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + # A distinct call_id: a second, separate request from the same key, not + # a fallback retry of the first. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_apply_to_key_alias_composes_with_scope_by_key_hash(time_controller, monkeypatch): + """Both listed keys are subject to the entry, but scope_by_key_hash + splits their buckets: exhausting one must not affect the other.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "daily", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_key_alias": ["key-a", "key-b"], + "scope_by_key_hash": True, + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a", api_key="hashA"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + # A distinct call_id: a second, separate request from the same key, not + # a fallback retry of the first. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a", api_key="hashA"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + # key-b is unaffected by key-a's exhausted bucket. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-b", api_key="hashB"), + cache=DualCache(), + data=_data(["end_user_id:u1"]), + call_type="completion", + ) + assert result is not None + + +# --------------------------------------------------------------------------- +# apply_to_models -- narrows which requested model an entry applies to, +# letting one entry rate-limit a whole fallback chain as a single unit +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_to_models_ignores_non_matching_model(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + for i in range(3): + data = {**_data(["end_user_id:u1"], call_id=f"call-{i}"), "model": "sonnet-chain"} + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion" + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_apply_to_models_enforces_for_the_listed_model(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data1 = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data1, call_type="completion") + data2 = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"} + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data2, call_type="completion") + + +@pytest.mark.asyncio +async def test_apply_to_models_shares_one_bucket_across_every_listed_model(time_controller, monkeypatch): + """The core "rate limit the whole chain" use case: a single limit shared + across every model named in apply_to_models, not one bucket per model.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain", "sonnet-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data1 = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data1, call_type="completion") + + data2 = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "sonnet-chain"} + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data2, call_type="completion") + + +@pytest.mark.asyncio +async def test_apply_to_models_composes_with_apply_to_key_alias(time_controller, monkeypatch): + """Both gates must pass -- the listed key requesting a non-listed model + is unaffected, and only the listed key requesting the listed model is + actually enforced.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [ + { + "name": "chain_cap", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + # premium-key requesting a non-listed model: apply_to_models alone must + # still exclude it, even though apply_to_key_alias matches. + for i in range(3): + data = {**_data(["end_user_id:u1"], call_id=f"wrong-model-{i}"), "model": "sonnet-chain"} + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), cache=DualCache(), data=data, call_type="completion" + ) + assert result is not None + + # A non-listed key requesting the listed model: apply_to_key_alias alone + # must still exclude it, even though apply_to_models matches. + for i in range(3): + data = {**_data(["end_user_id:u1"], call_id=f"wrong-key-{i}"), "model": "opus-chain"} + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="other-key"), cache=DualCache(), data=data, call_type="completion" + ) + assert result is not None + + # Both gates match: enforced. + data1 = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), cache=DualCache(), data=data1, call_type="completion" + ) + data2 = {**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"} + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), cache=DualCache(), data=data2, call_type="completion" + ) + + +@pytest.mark.asyncio +async def test_dollar_limit_respects_apply_to_models_at_accounting_time(time_controller, monkeypatch): + """The entry only applies to opus-chain; a non-listed model's spend must + not be charged against this bucket at all -- proves apply_to_models + gates async_log_success_event's tokens/dollars accounting, not just + admission.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "chain_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "sonnet-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": {"total_tokens": 0, "response_cost": 999.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # opus-chain was never charged -- still fully under its own limit. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"}, + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_apply_to_models_fallback_does_not_re_narrow_accounting_to_the_serving_model(time_controller, monkeypatch): + """ + Documented limitation, not a bug: apply_to_models is evaluated exactly + once, at admission, against the caller-requested model -- it is never + re-evaluated against whichever model a later fallback actually serves. + This request names "opus-chain" at admission (the entry applies), but its + response accounting reports "sonnet-chain" as the model that actually + served it, simulating Router falling back after opus-chain failed. The + spend must still land in the opus-chain-scoped bucket: the check already + ran and decided at admission, and is not re-run for the fallback target. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "chain_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data = {**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"} + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + + # The response accounting reports the fallback target as the model that + # actually served the request -- not the "opus-chain" admission decided. + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "model": "sonnet-chain", + "standard_logging_object": { + "total_tokens": 0, + "response_cost": 12.0, + "model": "sonnet-chain", + "model_group": "sonnet-chain", + }, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # The spend landed in the opus-chain-scoped bucket regardless -- a fresh + # opus-chain request is now over the limit. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"}, + call_type="completion", + ) + + +# --------------------------------------------------------------------------- +# _pre_call_with_fallbacks reruns admission for the same logical request: +# a repeat call_id must renew, not double-charge -- veria-ai finding on +# PR #36541 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_repeat_admission_for_the_same_call_id_and_key_renews_instead_of_double_charging( + time_controller, monkeypatch +): + """ + ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reruns the whole + pre-call pipeline (this hook included) once per fallback model on ANY + ProxyRateLimitError, not only one this hook itself raised, but keeps the + same litellm_call_id across every attempt (self.data is mutated in + place; only "model" changes). Without this fix, an unrelated rejection + (a different rate limiter, a budget cap) triggering N fallback attempts + would charge this hook's own "requests" cap N times for one logical + client call. A limit of 1 makes a double-charge directly observable: if + the second admission (same call_id, same key) charged again instead of + renewing, this would raise. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + key = _key(alias="key-a") + + await hook.async_pre_call_hook( + user_api_key_dict=key, cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + # Same call_id, same key, different model -- exactly what + # _pre_call_with_fallbacks produces for a fallback attempt of the same + # logical request. + result = await hook.async_pre_call_hook( + user_api_key_dict=key, + cache=DualCache(), + data={**_data(["end_user_id:u1"]), "model": "fallback-model"}, + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_a_forged_shared_call_id_from_a_different_key_does_not_get_a_free_renewal(time_controller, monkeypatch): + """ + Security regression: litellm_call_id is caller-controlled via the + x-litellm-call-id header (the same forgery vector + model_based_tag_rate_limits_hook's own pending-reservations mirror was + hardened against earlier in this PR). Two unrelated requests choosing + the identical call_id must not be able to renew each other's charge -- + only a second admission carrying the SAME authenticated key_hash as + whichever request first claimed that call_id may. A limit of 1 makes + this observable: if the second, different-key admission wrongly + renewed, it would succeed instead of raising. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-a", api_key="hashA"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="forged-call-id"), + call_type="completion", + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="key-b", api_key="hashB"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="forged-call-id"), + call_type="completion", + ) + + +# --------------------------------------------------------------------------- +# Concurrency: reservation at admission, release on success/failure/disconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_concurrency_limit_rejects_second_admission_until_release(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-1"), + call_type="completion", + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_on_success(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + async def one_request(call_id: str) -> None: + data = _data(["end_user_id:u1"], call_id=call_id) + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = {"litellm_call_id": call_id, "metadata": {"tags": ["end_user_id:u1"]}} + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + await one_request("call-1") + await asyncio.sleep(0) # let the fire-and-forget release task run + + # The slot was released, so a fresh request must be admitted again. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_nested_call_success_does_not_release_the_outer_calls_reservation(time_controller, monkeypatch): + """ + A nested LiteLLM call made inside the request (an LLM-judge guardrail, a + silent experiment) mints its own fresh litellm_call_id but inherits the + same ContextVar-held stash as the outer call, since it runs in the same + task rather than a separate one. The outer call's own concurrency + reservation must survive the nested call's admission and success + callback: it belongs to a different call id and must not be touched by + it. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-outer"), + call_type="completion", + ) + + # Nested call, same task, different tag and a fresh call id -- admits + # and completes entirely before the outer call's own success/failure + # callback ever fires. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u2"], call_id="call-nested"), + call_type="completion", + ) + await hook.async_log_success_event( + kwargs={"litellm_call_id": "call-nested", "metadata": {"tags": ["end_user_id:u2"]}}, + response_obj=None, + start_time=0, + end_time=0, + ) + await asyncio.sleep(0) + + # The outer call is still genuinely in flight -- its own reservation + # must still be held, so a second end_user_id:u1 request is rejected. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-outer-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_when_a_different_hook_rejects_the_request(time_controller, monkeypatch): + """ + model_based_tag_rate_limits_hook raises the identical ProxyRateLimitError + shape (detail["error"] == "tag_rate_limit_exceeded") this hook's own + admission does, since both hooks share the same rejection marker. + async_log_failure_event fires on every registered CustomLogger regardless + of which one raised, so this hook must still release its own successfully + reserved concurrency slot when the *other* hook is what rejected the + request -- skipping release just because the marker matches would leak + this hook's own slot until the safety TTL, even though nothing about this + hook's own admission failed. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + {"concurrency_limits": {"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]}}, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + + other_hooks_rejection = ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "type": "requests", "tag_id": "end_user_id"}, + headers={"retry-after": "60"}, + rate_limit_type=None, + model="gpt-4o", + llm_provider="litellm_proxy", + ) + kwargs = {"litellm_call_id": "call-1", "exception": other_hooks_rejection} + await hook.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # The slot was released despite the shared rejection marker, so a fresh + # request must be admitted again. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_concurrency_reservation_released_on_disconnect(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) + + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_concurrent_requests_do_not_share_each_others_reservation_state(time_controller, monkeypatch): + """Two logically distinct requests running as separate asyncio Tasks must + not see each other's pending-concurrency stash, even though both share + this hook instance -- the whole point of the ContextVar-based stash.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 5, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + async def one_request(call_id: str) -> int: + data = _data(["end_user_id:u1"], call_id=call_id) + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = {"litellm_call_id": call_id, "metadata": {"tags": ["end_user_id:u1"]}} + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + return 1 + + results = await asyncio.gather(one_request("call-a"), one_request("call-b")) + await asyncio.sleep(0) + assert results == [1, 1] + + +# --------------------------------------------------------------------------- +# Accounting: tokens/dollars via async_log_success_event +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dollar_limit_accounts_usage_and_rejects_once_over(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [{"name": "daily_spend", "tag_id": "end_user_id", "limit": 10.0, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": {"total_tokens": 0, "response_cost": 12.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key( + time_controller, monkeypatch +): + """ + kwargs at async_log_success_event time is Logging.model_call_details, not + the flat dict admission sees -- for a plain (non LITELLM_METADATA_ROUTES) + chat completion, kwargs["litellm_params"] carries a "litellm_metadata" key + that is always present but set to None, alongside the real, populated + "metadata" dict. get_metadata_variable_name_from_kwargs only checks key + presence, so it always resolved to "litellm_metadata" here and read no + tags/identity at all, silently dropping every token/dollar/key-hash/alias + accounting for this route shape. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [{"name": "daily_spend", "tag_id": "end_user_id", "limit": 10.0, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + kwargs = { + "litellm_call_id": "call-1", + "litellm_params": { + "litellm_metadata": None, + "metadata": {"tags": ["end_user_id:u1"], "user_api_key": "hash"}, + }, + "standard_logging_object": {"total_tokens": 0, "response_cost": 12.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_dollar_limit_respects_apply_to_key_alias_at_accounting_time(time_controller, monkeypatch): + """The entry only applies to `premium-key`; a non-listed key's spend must + not be charged against this bucket at all.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "daily_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_key_alias": ["premium-key"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="other-key"), cache=DualCache(), data=data, call_type="completion" + ) + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"], "user_api_key_alias": "other-key"}, + "standard_logging_object": {"total_tokens": 0, "response_cost": 999.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # premium-key was never charged -- still fully under its own limit. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(alias="premium-key"), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +# --------------------------------------------------------------------------- +# Config hot-reload: identity-based re-validation, no restart needed +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_config_reload_takes_effect_on_next_request(time_controller, monkeypatch): + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 100, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + + # Reload to a stricter config -- a fresh dict object, matching how a + # proxy config reload replaces litellm_settings.global_tag_rate_limits + # wholesale via setattr(litellm, key, value). A changed `limit` folds + # into the bucket's own policy fingerprint, so this is a fresh counter; + # the new, stricter limit=1 is still reachable in exactly one more call. + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-3"), + call_type="completion", + ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index e092a6d6272..f67989db777 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5496,6 +5496,65 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == primary_model + @pytest.mark.asyncio + async def test_cross_model_scoped_rejection_is_not_retried_via_fallback(self): + """ + veria-ai finding on PR #36541: an entry using ``apply_to_models`` to cap + an entire fallback chain as one unit is defeated by this exact mechanism + if a fallback model isn't also listed in ``apply_to_models`` -- the + rejection here is a deliberate "this whole chain is capped" decision, + not a "this one model is unhealthy" signal, so retrying against an + unlisted fallback silently serves a request the operator's policy meant + to block. ``detail["cross_model_scope"]`` is the marker + global_tag_rate_limits_hook sets for exactly this case; the fallback + handler must re-raise immediately instead of trying any fallback model. + """ + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + primary_model = "opus-chain" + + processor = ProxyBaseLLMRequestProcessing(data={"model": primary_model}) + + call_count = 0 + + async def mock_pre_call_logic(**kwargs): + nonlocal call_count + call_count += 1 + raise ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "cross_model_scope": True}, + headers={"retry-after": "30"}, + ) + + mock_router = MagicMock() + mock_router.fallbacks = [{"opus-chain": ["sonnet-chain"]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + with pytest.raises(ProxyRateLimitError): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert call_count == 1 + assert processor.data["model"] == primary_model + @pytest.mark.asyncio async def test_real_parallel_request_limiter_model_tpm_limit_triggers_fallback(self): """ From 33b19104c50e6ba8564cf976dc925585fc31a21e Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:26:38 -0400 Subject: [PATCH 38/44] fix(rate-limiting): fix cross-attempt accounting loss, trim comment verbosity Track every model an admission attempt saw for a call_id (admitted_models) instead of overwriting a single field, so an apply_to_models-scoped entry that matched an earlier _pre_call_with_fallbacks attempt still gets its success-time token/dollar accounting even after a later attempt admits with a different, out-of-scope model. Also trims the module docstring, stash dataclass, and several inline comments down to essential rationale per CLAUDE.md's comment policy, and corrects the docstring's claim that the Logging object doesn't exist yet at admission time (it does; the real reason for the ContextVar-based stash is that _pre_call_with_fallbacks builds a fresh one per retry). --- .../hooks/global_tag_rate_limits_hook.py | 244 +++++++----------- .../hooks/test_global_tag_rate_limits_hook.py | 113 ++++++++ 2 files changed, 203 insertions(+), 154 deletions(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index ce8f59461d0..d48ae8fa643 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -1,70 +1,37 @@ """ Tag-scoped token, request, dollar, and concurrency rate limits declared once, -globally, in `litellm_settings.global_tag_rate_limits` -- enforced once per -request in `async_pre_call_hook`, before Router does any routing, so a limit -applies regardless of which model or fallback chain the request ends up -hitting. +globally, in `litellm_settings.global_tag_rate_limits` and enforced in +`async_pre_call_hook`, before Router does any routing -- so a limit applies +regardless of which model or fallback chain the request ends up hitting. -This is the model-independent sibling of `model_based_tag_rate_limits_hook`, -which enforces the same `TagRateLimitEntry` shape but nested per-deployment -under `model_info.tag_rate_limits`, once per routing hop -(`async_filter_deployments`). A global entry has no deployment/routing-group -to reconcile -- there is exactly one config value, read once -- so this hook -reuses that sibling's free, already-hardened helper functions -(`entry_applies`, the Lua atomic check-and-increment scripts, cache -partitioning, bucket-key hashing primitives) from `tag_rate_limits_shared.py` -rather than duplicating them, but implements its own, much smaller -admission/accounting engine: no `_LimitsIndex`, no routing-group or -team-alias resolution, no per-deployment dedup signatures. +Model-independent sibling of `model_based_tag_rate_limits_hook`, which +enforces the same `TagRateLimitEntry` shape per-deployment instead. Reuses +that sibling's shared helpers (`entry_applies`, the atomic Lua scripts, cache +partitioning, bucket-key hashing) from `tag_rate_limits_shared.py`, but has +its own smaller admission/accounting engine with no routing-group or +per-deployment concerns. -Three independent entry-level knobs decide who a global entry applies to and -how its bucket is shared: +Three entry-level knobs: `apply_to_key_alias` scopes an entry to specific +virtual-key aliases; `apply_to_models` scopes it to specific caller-facing +model names, letting one entry cap a whole fallback chain as a unit (a +rejection then carries `detail["cross_model_scope"] = True` so +`_pre_call_with_fallbacks` re-raises instead of silently admitting the +request through an unlisted fallback model); `scope_by_key_hash` controls +whether matching keys share one bucket or each gets its own. -- `apply_to_key_alias`: unset means every request, any key, any model. Set - to a list of virtual-key aliases, only those keys' requests count. -- `apply_to_models`: unset means every model. Set to a list of model names, - only requests whose caller-facing `model` field is in that list count -- - letting one entry rate-limit a whole fallback chain as a single unit by - naming every model in the chain. Each check is a fresh, independent - evaluation of `_entry_applies` against whatever `model` is current at that - moment, not a one-time decision that then sticks for the rest of the - request. Two concrete consequences follow from that: - (1) if the request's own model fails mid-flight and Router internally - retries a different model for the *same* admitted call, that retry is - never re-checked -- the original admission (against the originally - requested model) already stands, so an operator who needs the limit to - track whichever model actually ends up serving a request needs - `model_info.tag_rate_limits` instead; but - (2) if this hook's own admission *rejects* the request, - `common_request_processing.py` would otherwise catch that rejection and - retry the whole pre-call pipeline against - `litellm_settings.fallbacks`/`router_settings.fallbacks`, with - `data["model"]` mutated to the fallback target -- silently admitting the - request via a model outside `apply_to_models`, defeating the cap. A - rejection from an `apply_to_models`-scoped entry carries - `detail["cross_model_scope"] = True` for exactly this reason: - `_pre_call_with_fallbacks` checks that marker and re-raises immediately - instead of trying any fallback, so this bypass is closed regardless of - whether the fallback chain is also listed in `apply_to_models`. -- `scope_by_key_hash` (already exists on `TagRateLimitEntry`): whether the - keys an entry applies to share one bucket, or each gets its own. - -`async_pre_call_hook` runs before Router constructs `Logging`/`litellm_logging_obj` -for this request (see `common_request_processing.py`: `pre_call_hook` fires -well before `base_process_llm_request` builds the logging object), so unlike -`model_based_tag_rate_limits_hook` this hook cannot stash pending concurrency -reservations on `data["litellm_logging_obj"].model_call_details` -- that -object doesn't exist yet. Per-request state is instead kept on a -`ContextVar`-based stash, the same established pattern -`parallel_request_limiter_v3.py`'s v3 handler already uses for exactly this -problem, with one difference: the stash here is a dict keyed by -`litellm_call_id` rather than one shared mutable instance with an -overwritable "owner" field, so a nested LiteLLM call made inside the request -(e.g. a guardrail's own LLM judge call) -- which mints its own fresh call id -but inherits the same ContextVar-held ancestor context, not a separate one --- gets its own isolated entry instead of overwriting the outer call's and -having its own success callback release the outer call's still-pending -reservations early. +`data["litellm_logging_obj"]` does exist by the time `async_pre_call_hook` +runs, but `_pre_call_with_fallbacks` re-runs the whole pre-call pipeline +(building a fresh `Logging` object each time) once per fallback model on the +same `litellm_call_id`, so unlike `model_based_tag_rate_limits_hook`'s +per-Router-hop reservations, a stash keyed to one attempt's own +`model_call_details` wouldn't survive to a later attempt. Per-request state +instead lives on a `ContextVar`-based stash (the same pattern +`parallel_request_limiter_v3.py` uses for the identical admission-to-release +problem), keyed by `litellm_call_id` so a nested LiteLLM call sharing the +same inherited context (an LLM-judge guardrail, for example) gets its own +isolated entry instead of releasing the outer call's still-pending +reservation early. Confirmed live: a real streaming client disconnect +correctly releases its concurrency reservation through this mechanism. """ import asyncio @@ -84,7 +51,7 @@ from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_f from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( - _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching model_based_tag_rate_limits_hook's identical import + _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # shared private helper, reused by model_based_tag_rate_limits_hook too ) from litellm.proxy.hooks.tag_rate_limits_shared import ( ATOMIC_UNITS as _ATOMIC_UNITS, @@ -146,7 +113,7 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( ) from litellm.proxy.utils import InternalUsageCache from litellm.router_strategy.tag_based_routing import ( - _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching model_based_tag_rate_limits_hook's identical import + _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # shared private helper, reused by model_based_tag_rate_limits_hook too ) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import TagRateLimitEntry, TagRateLimits @@ -160,16 +127,23 @@ else: Span: TypeAlias = object +def _entry_applies_any_admitted_model( + entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, admitted_models: frozenset[str] +) -> bool: + """Same as `_entry_applies`, except an `apply_to_models`-scoped entry + counts as applying if ANY model an admission attempt for this call_id + saw was in scope -- not just whichever model the call ultimately served. + A `_pre_call_with_fallbacks` retry re-admits with a different model for + the same call_id, and an entry that matched an earlier attempt must + still get its success-time accounting.""" + if not admitted_models: + return _entry_applies(entry, tags, key_alias, None) + return any(_entry_applies(entry, tags, key_alias, model) for model in admitted_models) + + def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: - """ - Global-hook equivalent of `model_based_tag_rate_limits_hook._hash_tag`, - without a `model_group`/deployment-scope/team-scope dimension -- a global - entry has none of those. Namespaced under `tag_rl:global:` so it can never - collide with that sibling hook's own `tag_rl:{model_group}:...` keys even - if an operator names a deployment "global": every key also differs by - `unit`/`name`/`tag_id`/`_policy_fingerprint`, and the two hooks' entries - are never meant to share a bucket in the first place. - """ + """Namespaced under `tag_rl:global:` so it never collides with + `model_based_tag_rate_limits_hook`'s own `tag_rl:{model_group}:...` keys.""" key_suffix: Final = f":key:{key_hash}" if key_hash is not None else "" policy_suffix: Final = f":policy:{_policy_fingerprint(entry)}" return f"tag_rl:global:{unit}:{entry.name}:{entry.tag_id}:{_fixed_length_identity(tag_value)}{key_suffix}{policy_suffix}" @@ -203,45 +177,30 @@ class _CachePartition: @dataclass(slots=True) class _GlobalTagRateLimitStash: """Per-call bookkeeping `async_pre_call_hook` hands to that same call's - success/failure/disconnect callbacks -- see module docstring for why - this lives on a `ContextVar`, not `model_call_details`. + success/failure/disconnect callbacks -- see module docstring for why this + lives on a `ContextVar`, not `model_call_details`. - Keyed by `litellm_call_id` in the dict below rather than one shared - mutable instance with an overwritable "owner" field: a nested LiteLLM - call made inside the request (an LLM-judge guardrail, a silent - experiment) that mints its own fresh call id runs inside the *same* - inherited context, not a separate one, so a single shared instance's - owner field would get reassigned to the nested call and its own - success callback would then release the outer call's still-pending - reservations early -- letting extra same-tag requests through while - the outer request is still genuinely in flight. Keying by call id - isolates each call's own reservations regardless of nesting. + Keyed by `litellm_call_id` rather than one shared mutable instance so a + nested LiteLLM call (an LLM-judge guardrail) that mints its own call id + but inherits the same context doesn't release the outer call's + still-pending reservation early. """ admission_time: float | None = None - # The caller-facing `model` admission read from `data.get("model")`, so - # async_log_success_event's tokens/dollars accounting gates - # apply_to_models against the same, originally-requested model admission - # decided on -- not whatever model a later fallback actually served. - model: str | None = None + # Every model any admission attempt for this call_id has classified + # entries against, accumulated rather than overwritten: a + # _pre_call_with_fallbacks retry re-runs admission with a *different* + # model for the same call_id, and an apply_to_models entry that matched + # an earlier attempt must still get its accounting at success time even + # though the request ultimately serves from a later attempt's model. + admitted_models: frozenset[str] = field(default_factory=frozenset) pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue - # "requests" keys already charged for this call_id -- veria-ai finding: - # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reruns the whole - # pre-call pipeline (this hook included) once per fallback model on ANY - # ProxyRateLimitError, not only one this hook itself raised, but reuses - # the same litellm_call_id (self.data is mutated in place, only `model` - # changes) across every attempt -- so this stash is the SAME object each - # time. A "requests" check matching an already-charged key here renews - # at zero net cost instead of charging a second unit for the same - # logical request; see async_pre_call_hook's own comment for how. + # Keys already charged for this call_id, so a fallback retry (same + # litellm_call_id, different model) renews instead of double-charging. charged_request_keys: list[str] = field(default_factory=list) # mutable-ok: see comment above - # The server-authenticated key_hash (UserAPIKeyAuth.api_key) of whichever - # call first claimed this stash. litellm_call_id is caller-controlled via - # the x-litellm-call-id header (the exact forgery vector - # model_based_tag_rate_limits_hook's own pending-reservations mirror was - # hardened against earlier), so two unrelated requests sharing a - # caller-chosen id must not be allowed to "renew" each other's charge -- - # only a later admission carrying this same, authenticated key_hash may. + # key_hash of whoever first claimed this stash. litellm_call_id is + # caller-controlled (x-litellm-call-id), so only a later admission with + # the same authenticated key_hash may renew this stash's charges. owner_key_hash: str | None = None @@ -392,12 +351,8 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], ) -> tuple[int | None, tuple[float, ...]]: - """All-or-nothing atomic admission across `checks` -- see - `model_based_tag_rate_limits_hook._PROXY_ModelBasedTagRateLimitsHook._atomic_check_and_increment`'s - own docstring for the full rationale (refund-on-rollback, why a - raising key's own outcome is never refunded); identical logic, - duplicated rather than shared since it lives as instance methods - rather than free functions.""" + """All-or-nothing atomic admission across `checks`: on a rejection, + refunds every check admitted earlier in this batch.""" if not checks: return None, () admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection @@ -563,12 +518,10 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o if config is None: return data - # async_pre_call_hook fires once per request in the common case, but - # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks can re-run - # this same pipeline once per fallback model on any ProxyRateLimitError - # (not only one this hook raised) -- see charged_request_keys' own - # docstring for how a repeat run for the same call_id renews rather - # than re-charges both "requests" and "concurrency" checks below. + # _pre_call_with_fallbacks can re-run this pipeline once per fallback + # model on any ProxyRateLimitError, reusing the same litellm_call_id -- + # see charged_request_keys for how a repeat run renews instead of + # re-charging. stash: Final = _claim_stash_for_data(data) metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(data) @@ -577,17 +530,16 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o key_hash: Final = user_api_key_dict.api_key model: Final = data.get("model") if isinstance(data.get("model"), str) else None - # Only a repeat admission carrying the SAME authenticated key_hash as - # whichever call first claimed this stash may renew its charges -- - # see owner_key_hash's own docstring for why a bare call_id match is - # not enough. First admission for this stash claims ownership here. + # First admission for this stash claims ownership; only a later one + # with the same key_hash may renew its charges (see owner_key_hash). if stash.owner_key_hash is None: stash.owner_key_hash = key_hash renewal_allowed: Final = stash.owner_key_hash == key_hash now: Final = self._time_provider().timestamp() stash.admission_time = now - stash.model = model + if renewal_allowed and model is not None: + stash.admitted_models = stash.admitted_models | frozenset((model,)) classified: Final = self._classify(config, tags, key_alias, key_hash, now, model) if not classified: return data @@ -615,15 +567,8 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o check.key, check.entry.limit, # A key already charged/reserved for this call_id (an - # earlier _pre_call_with_fallbacks attempt for the - # same logical request) renews at zero net cost - # instead of charging or reserving a second unit -- - # folded into this same all-or-nothing batch so a - # rollback here (some other check in the batch - # rejecting) refunds that zero-cost renewal as a - # genuine no-op, same reasoning as - # model_based_tag_rate_limits_hook's identical fix - # for its own per-hop retries. + # earlier fallback attempt for the same request) renews + # at zero net cost instead of charging a second unit. 0.0 if renewal_allowed and ( @@ -642,11 +587,9 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o failing_check.unit, failing_check.entry, failing_check.tag_value, model, current=values[0] ) - # Only genuinely new reservations, never a key already in - # already_reserved_concurrency_keys: that key's own check just - # renewed at zero net cost above, so re-adding it here would - # make release (which decrements once per queued entry) decrement - # twice for a counter that was only ever incremented once. + # Exclude already_reserved_concurrency_keys: that key renewed at + # zero cost above, so re-adding it would make release decrement + # twice for a counter only ever incremented once. concurrency_reservations: Final = tuple( (check.key, _partition_key(check.entry)) for check in atomic_checks @@ -655,12 +598,9 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o if concurrency_reservations: stash.pending_concurrency_keys.extend(concurrency_reservations) # mutable-ok: see field's own docstring - # Only recorded when renewal_allowed: an admission that didn't - # own this stash (a call_id collision from a different key_hash) - # must not contaminate the rightful owner's own renewal - # tracking, or a later, genuine fallback retry from the owner - # could wrongly treat the impostor's charge as its own and - # renew for free. + # Only recorded when renewal_allowed, so a call_id collision from + # a different key_hash can't contaminate the rightful owner's + # renewal tracking. request_keys: Final = ( tuple( check.key @@ -684,15 +624,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o await self._release_keys(release_keys) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - # No special-case skip for this hook's own tag_rate_limit_exceeded - # rejection: that rejection never reaches the point where a - # concurrency reservation is queued (see async_pre_call_hook), so - # stash.pending_concurrency_keys is already empty in that case and - # the check below naturally no-ops. Skipping release based on the - # exception's error marker alone would be wrong here, since - # model_based_tag_rate_limits_hook raises the identical marker -- - # that rejection can land after this hook already reserved a slot - # for this same request, and that slot must still be released. + # Always release regardless of which hook raised: this hook's own + # rejection never reserves a slot, so pending_concurrency_keys is + # already empty in that case and the check below no-ops; a rejection + # from model_based_tag_rate_limits_hook (same error marker) can still + # land after this hook already reserved its own slot. stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) if stash is None or not stash.pending_concurrency_keys: return @@ -735,7 +671,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o if stash is not None and stash.admission_time is not None else self._time_provider().timestamp() ) - model: Final = stash.model if stash is not None else None + admitted_models: Final = stash.admitted_models if stash is not None else frozenset() increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType( { "tokens": float(standard_logging_object.get("total_tokens") or 0), @@ -752,7 +688,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o tag_value = _extract_identity(tags, entry.tag_id) if tag_value is None: continue - if not _entry_applies(entry, tags, key_alias, model): + if not _entry_applies_any_admitted_model(entry, tags, key_alias, admitted_models): continue increment_value = increment_by_unit[unit] if increment_value == 0: diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 428037986e5..e48a1278cd7 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -517,6 +517,79 @@ async def test_apply_to_models_fallback_does_not_re_narrow_accounting_to_the_ser ) +@pytest.mark.asyncio +async def test_apply_to_models_accounts_when_a_fallback_retry_re_admits_with_a_different_model( + time_controller, monkeypatch +): + """ + Unlike the test above (one admission, a later Router-level fallback + reported only at success time), this simulates + _pre_call_with_fallbacks itself re-running async_pre_call_hook for the + SAME call_id with a DIFFERENT model after some other hook rejected the + original one. The first admission (opus-chain, in apply_to_models scope) + must still get its success-time accounting even though the second + admission (sonnet-chain, out of scope) is the one that actually proceeds + -- overwriting a single "last admitted model" field would silently drop + the spend for the in-scope entry. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [ + { + "name": "chain_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_models": ["opus-chain"], + } + ] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"}, + call_type="completion", + ) + # _pre_call_with_fallbacks retry: same call_id, fallback model outside apply_to_models. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "sonnet-chain"}, + call_type="completion", + ) + + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "model": "sonnet-chain", + "standard_logging_object": { + "total_tokens": 0, + "response_cost": 12.0, + "model": "sonnet-chain", + "model_group": "sonnet-chain", + }, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # The $12 spend landed in the opus-chain-scoped bucket, so a fresh + # opus-chain request is now over the $10 limit and rejected. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-2"), "model": "opus-chain"}, + call_type="completion", + ) + + # --------------------------------------------------------------------------- # _pre_call_with_fallbacks reruns admission for the same logical request: # a repeat call_id must renew, not double-charge -- veria-ai finding on @@ -798,6 +871,46 @@ async def test_concurrency_reservation_released_on_disconnect(time_controller, m assert result is not None +@pytest.mark.asyncio +async def test_concurrency_reservation_released_when_disconnect_runs_in_a_forked_task(time_controller, monkeypatch): + """ + common_request_processing.py's real disconnect path runs the release call + inside a task forked via asyncio.create_task AFTER admission already ran + in the parent task (_await_llm_call_cancelling_on_disconnect's `monitor` + task, the streaming generator's own task), not in the same coroutine as + admission the way the test above calls it. A ContextVar-based stash that + only happens to work when both calls share one coroutine has caused a + real leak in this codebase before, so this exercises the actual + parent-then-forked-child topology instead. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + + async def release_in_child_task() -> None: + await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) + + await asyncio.create_task(release_in_child_task()) + + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + @pytest.mark.asyncio async def test_concurrent_requests_do_not_share_each_others_reservation_state(time_controller, monkeypatch): """Two logically distinct requests running as separate asyncio Tasks must From 2130f7c600288d8b56acc5a150ca3df7d6ef28a3 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 11:51:14 -0400 Subject: [PATCH 39/44] fix(rate-limiting): close three gaps in the global tag hook and fallback guard - Wire order_tags_for_identity_resolution into global_tag_rate_limits_hook's admission and success-event tag resolution, matching the sibling hook. Without it a caller could put a forged tag ahead of the policy-backed inherited one and dodge or mis-bucket every global limit. Confirmed exploitable pre-fix, blocked post-fix, via a direct adversarial repro. - Only record a model into the stash's admitted-models set once an admission attempt clears every check, not unconditionally at the top -- a rejected attempt's model could otherwise still drive a later successful attempt's token/dollar accounting for an apply_to_models entry that never actually admitted the request. - Check cross_model_scope on a fallback attempt's own rejection inside _pre_call_with_fallbacks's retry loop, not only on the original exception before the loop starts -- a chain-wide apply_to_models cap covering the first fallback too was previously bypassable by a second, uncovered fallback model. Each fix has a regression test confirmed to fail on the pre-fix code and pass on the fix. --- litellm/proxy/common_request_processing.py | 14 +- .../hooks/global_tag_rate_limits_hook.py | 29 ++- .../hooks/test_global_tag_rate_limits_hook.py | 186 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 64 ++++++ 4 files changed, 288 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0f54b50a16f..a687aaf2c5c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2040,7 +2040,19 @@ class ProxyBaseLLMRequestProcessing: route_type=route_type, llm_router=llm_router, ) - except ProxyRateLimitError: + except ProxyRateLimitError as fallback_exc: + # A fallback attempt's own rejection can carry the + # identical cross_model_scope marker (this fallback + # model is itself covered by the same apply_to_models + # chain-wide cap) -- continuing to the next fallback + # would silently serve the request through a model + # outside that cap, defeating it just as much as not + # checking the original exception would. + if ( + isinstance(fallback_exc.detail, Mapping) + and fallback_exc.detail.get("cross_model_scope") is True + ): + raise continue except BaseException: self.data["model"] = original_model diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index d48ae8fa643..92b905db49f 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -102,6 +102,9 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( from litellm.proxy.hooks.tag_rate_limits_shared import ( fixed_length_identity as _fixed_length_identity, ) +from litellm.proxy.hooks.tag_rate_limits_shared import ( + order_tags_for_identity_resolution as _order_tags_for_identity_resolution, +) from litellm.proxy.hooks.tag_rate_limits_shared import ( partition_key as _partition_key, ) @@ -391,6 +394,16 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o "global_tag_rate_limits_hook: failed to release concurrency slot %s: %s", key, e ) + @staticmethod + def _record_admitted_model(stash: _GlobalTagRateLimitStash, model: str | None, renewal_allowed: bool) -> None: + """Only called once this admission attempt has cleared every check + without raising -- a rejected attempt's model must never join + admitted_models, or a later successful attempt's accounting could + wrongly credit an apply_to_models entry that never actually admitted + this request under that model.""" + if renewal_allowed and model is not None: + stash.admitted_models = stash.admitted_models | frozenset((model,)) + @staticmethod def _ttl_for(unit: _LimitUnit, entry: TagRateLimitEntry) -> int: if unit == "concurrency": @@ -525,7 +538,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o stash: Final = _claim_stash_for_data(data) metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(data) - tags: Final = _get_tags_from_request_kwargs(data, metadata_variable_name=metadata_variable_name) + tags: Final = _order_tags_for_identity_resolution( + _get_tags_from_request_kwargs(data, metadata_variable_name=metadata_variable_name), + data, + metadata_variable_name, + ) key_alias: Final = user_api_key_dict.key_alias key_hash: Final = user_api_key_dict.api_key model: Final = data.get("model") if isinstance(data.get("model"), str) else None @@ -538,10 +555,9 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o now: Final = self._time_provider().timestamp() stash.admission_time = now - if renewal_allowed and model is not None: - stash.admitted_models = stash.admitted_models | frozenset((model,)) classified: Final = self._classify(config, tags, key_alias, key_hash, now, model) if not classified: + self._record_admitted_model(stash, model, renewal_allowed) return data read_only_checks: Final = tuple(c for c in classified if not c.is_atomic) @@ -613,6 +629,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o if request_keys: stash.charged_request_keys.extend(request_keys) # mutable-ok: see field's own docstring + self._record_admitted_model(stash, model, renewal_allowed) return data async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: @@ -662,7 +679,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name) - tags: Final = _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name) + tags: Final = _order_tags_for_identity_resolution( + _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name), + kwargs, + metadata_variable_name, + ) if not tags: return diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index e48a1278cd7..97f1e065c68 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -590,6 +590,101 @@ async def test_apply_to_models_accounts_when_a_fallback_retry_re_admits_with_a_d ) +@pytest.mark.asyncio +async def test_a_rejected_admission_attempts_model_does_not_drive_later_accounting(time_controller, monkeypatch): + """ + A fallback retry's FIRST attempt can itself be rejected (by this same + entry, or a different hook) before it ever admits. That rejected + attempt's model must not join the stash's admitted-models history: a + later, successful attempt against a different (out-of-scope) model must + not have its accounting wrongly credited to an apply_to_models entry + that never actually admitted this request. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [ + { + "name": "conc-a", + "tag_id": "end_user_id", + "limit": 1, + "period_seconds": 60, + "apply_to_models": ["model-a"], + } + ] + }, + "dollar_limits": { + "limits": [ + { + "name": "chain_spend", + "tag_id": "end_user_id", + "limit": 10.0, + "period_seconds": 86400, + "apply_to_models": ["model-a"], + } + ] + }, + }, + ) + hook = _make_hook(time_controller) + + # Occupy model-a's only concurrency slot with an unrelated call. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="occupier"), "model": "model-a"}, + call_type="completion", + ) + + # call-1 attempt #1: model-a, rejected (slot taken) -- never truly admitted. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-a"}, + call_type="completion", + ) + # call-1 attempt #2 (fallback retry): model-b, not in apply_to_models=[model-a], admits. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-b"}, + call_type="completion", + ) + + kwargs = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + "model": "model-b", + "standard_logging_object": {"total_tokens": 0, "response_cost": 50.0, "model": "model-b", "model_group": "model-b"}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # Release the occupier's concurrency slot so the final check below is + # gated only by chain_spend (dollars), not by conc-a still being full. + await hook.async_log_success_event( + kwargs={"litellm_call_id": "occupier", "metadata": {"tags": ["end_user_id:u1"]}}, + response_obj=None, + start_time=0, + end_time=0, + ) + await asyncio.sleep(0) + + # chain_spend (apply_to_models=[model-a]) must still be empty: model-a's + # own admission attempt was rejected, never admitted, so a fresh + # model-a request is still allowed under the $100 limit. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-2"), "model": "model-a"}, + call_type="completion", + ) + assert result is not None + + # --------------------------------------------------------------------------- # _pre_call_with_fallbacks reruns admission for the same logical request: # a repeat call_id must renew, not double-charge -- veria-ai finding on @@ -1116,3 +1211,94 @@ async def test_config_reload_takes_effect_on_next_request(time_controller, monke data=_data(["end_user_id:u1"], call_id="call-3"), call_type="completion", ) + + +# --------------------------------------------------------------------------- +# Identity resolution: policy-backed tags must win over caller-supplied ones +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_caller_supplied_tag_cannot_shadow_the_policy_backed_identity_tag(time_controller, monkeypatch): + """ + _merge_tags (litellm_pre_call_utils.py) keeps caller-supplied tags first + in the merged tags list, appending key/team/project-contributed tags only + if not already present. Since extract_identity/entry_applies resolve a + tag_id by first-match-by-prefix, an authenticated caller could otherwise + submit e.g. company_id:attacker-chosen ahead of the key's real + company_id:real-company (surfaced via metadata.inherited_tags) and have + every company_id-scoped entry resolve to the forged value instead of the + real one -- letting the caller dodge the limit entirely by rotating + fabricated identities, or evade being charged against their own real + bucket. The hook must order tags so inherited_tags wins. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "per-company", "tag_id": "company_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + key = _key() + + poisoned_data = { + "litellm_call_id": "attack-1", + "metadata": { + "tags": ["company_id:attacker-chosen", "company_id:real-company"], + "inherited_tags": ["company_id:real-company"], + }, + } + await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=poisoned_data, call_type="completion") + + # The real company's own bucket must have been charged by the attack + # request, not a bucket keyed to the attacker's forged value -- so a + # second, genuine company_id:real-company request is now rejected. + victim_data = { + "litellm_call_id": "victim-1", + "metadata": {"tags": ["company_id:real-company"], "inherited_tags": ["company_id:real-company"]}, + } + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=victim_data, call_type="completion") + + +@pytest.mark.asyncio +async def test_success_accounting_also_resolves_identity_from_the_policy_backed_tag(time_controller, monkeypatch): + """Same forged-tag scenario as the admission-time test above, but for + async_log_success_event's own identity resolution (tokens/dollars).""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [{"name": "per-company-spend", "tag_id": "company_id", "limit": 10.0, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + kwargs = { + "litellm_call_id": "attack-1", + "metadata": { + "tags": ["company_id:attacker-chosen", "company_id:real-company"], + "inherited_tags": ["company_id:real-company"], + }, + "standard_logging_object": {"total_tokens": 0, "response_cost": 20.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + # The $20 spend must have landed against company_id:real-company, so a + # fresh request under the genuine identity is now over the $10 limit. + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={ + "litellm_call_id": "victim-1", + "metadata": {"tags": ["company_id:real-company"], "inherited_tags": ["company_id:real-company"]}, + }, + call_type="completion", + ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f67989db777..805ee88a136 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5555,6 +5555,70 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert call_count == 1 assert processor.data["model"] == primary_model + @pytest.mark.asyncio + async def test_cross_model_scoped_rejection_mid_chain_stops_further_fallback_attempts(self): + """ + Bugbot finding: the original exception is checked for + cross_model_scope before the fallback loop starts, but a LATER + fallback attempt's own rejection was never checked the same way -- + the loop's `except ProxyRateLimitError: continue` swallowed it and + moved on to the next fallback model. If a chain-wide apply_to_models + cap covers both the primary model and the first fallback, and a + second fallback model isn't covered, this let the second fallback + silently serve the request the cap was meant to block. The original + (non-scoped) rejection enters the loop normally; the FIRST fallback's + own rejection carries cross_model_scope=True and must stop the loop + immediately, never reaching the second fallback. + """ + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + primary_model = "opus-chain" + + processor = ProxyBaseLLMRequestProcessing(data={"model": primary_model}) + + attempted_models: list[str] = [] + + async def mock_pre_call_logic(**kwargs): + attempted_models.append(processor.data["model"]) + if processor.data["model"] == primary_model: + # Original attempt: a plain, non-scoped rejection (e.g. a + # per-deployment limit), not the chain-wide cap itself. + raise ProxyRateLimitError(detail={"error": "tag_rate_limit_exceeded"}, headers={"retry-after": "30"}) + if processor.data["model"] == "sonnet-chain": + # First fallback: rejected by the SAME chain-wide cap. + raise ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "cross_model_scope": True}, + headers={"retry-after": "30"}, + ) + raise AssertionError(f"must not attempt a second fallback model: {processor.data['model']}") + + mock_router = MagicMock() + mock_router.fallbacks = [{"opus-chain": ["sonnet-chain", "haiku-chain"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + with pytest.raises(ProxyRateLimitError) as exc_info: + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert attempted_models == [primary_model, "sonnet-chain"] + assert exc_info.value.detail.get("cross_model_scope") is True + assert processor.data["model"] == primary_model + @pytest.mark.asyncio async def test_real_parallel_request_limiter_model_tpm_limit_triggers_fallback(self): """ From 7db9895c681d54e878a13b57493feae58e0ef0de Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 26 Aug 2026 13:23:56 -0400 Subject: [PATCH 40/44] fix(rate-limiting): fix success-event identity ordering, add failure-hook release, drop tag disclosure - Pass the correctly-nested litellm_params object (not raw model_call_details kwargs) to order_tags_for_identity_resolution in async_log_success_event. kwargs has no top-level metadata at that point, so the previous call was a silent no-op: a forged caller tag still won at accounting time even after the admission-time fix. Confirmed via a direct repro against the real Logging pipeline, and rewrote the existing regression test's kwargs shape to match production instead of a shape that happened to hide the bug. - Add async_post_call_failure_hook, releasing a reservation from an earlier admission attempt when _pre_call_with_fallbacks exhausts every fallback and re-raises without ever running the real LLM call -- the only other release paths (success/failure/disconnect) are tied to that call, which never happens. - Drop tag_value from the client-facing rejection detail: it can resolve from inherited_tags (server-assigned key/team/project metadata), and echoing it back would disclose that identity to the caller who got rejected. Each fix has a regression test confirmed to fail on the pre-fix code. --- .../hooks/global_tag_rate_limits_hook.py | 42 +++++-- .../hooks/test_global_tag_rate_limits_hook.py | 104 +++++++++++++++++- 2 files changed, 133 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index 92b905db49f..f0a3c481ab6 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -504,7 +504,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o "error": "tag_rate_limit_exceeded", "type": unit, "tag_id": entry.tag_id, - "tag_value": tag_value, + # tag_value deliberately excluded: it can resolve from + # inherited_tags (server-assigned key/team/project metadata), + # and echoing it back would disclose that identity to the + # caller. verbose_proxy_logger.debug above still logs it + # server-side for observability. "limit_name": entry.name, "limit": entry.limit, "period_seconds": entry.period_seconds, @@ -632,26 +636,44 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o self._record_admitted_model(stash, model, renewal_allowed) return data - async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: - stash: Final = _stash_for_call(_call_id_from_kwargs(request_data)) + async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None: + stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs)) if stash is None or not stash.pending_concurrency_keys: return release_keys: Final = tuple(stash.pending_concurrency_keys) stash.pending_concurrency_keys.clear() await self._release_keys(release_keys) + async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: + await self._release_pending_for_call_id(request_data) + + async def async_post_call_failure_hook( + self, + request_data: dict, # mutable-ok: must match CustomLogger.async_post_call_failure_hook's own base signature exactly + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + """ + A request that never reaches Router (every fallback model also + rejected, or none configured) never runs the actual LLM call, so + neither async_log_success_event nor async_log_failure_event -- both + tied to that call's own wrapper -- ever fires for it. This is the + only remaining release path for a reservation from an earlier, + successful admission attempt in the same _pre_call_with_fallbacks + chain. litellm_call_id survives proxy/utils.py's own stripping here + (only litellm_logging_obj is popped), so the same ContextVar-based + stash lookup as the other release hooks still works. + """ + await self._release_pending_for_call_id(request_data) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: # Always release regardless of which hook raised: this hook's own # rejection never reserves a slot, so pending_concurrency_keys is # already empty in that case and the check below no-ops; a rejection # from model_based_tag_rate_limits_hook (same error marker) can still # land after this hook already reserved its own slot. - stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) - if stash is None or not stash.pending_concurrency_keys: - return - release_keys: Final = tuple(stash.pending_concurrency_keys) - stash.pending_concurrency_keys.clear() - await self._release_keys(release_keys) + await self._release_pending_for_call_id(kwargs) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) @@ -681,7 +703,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o tags: Final = _order_tags_for_identity_resolution( _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name), - kwargs, + litellm_params_for_metadata, metadata_variable_name, ) if not tags: diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 97f1e065c68..c77cde23484 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -966,6 +966,54 @@ async def test_concurrency_reservation_released_on_disconnect(time_controller, m assert result is not None +@pytest.mark.asyncio +async def test_concurrency_reservation_released_when_every_fallback_is_exhausted(time_controller, monkeypatch): + """ + When _pre_call_with_fallbacks exhausts every fallback model (another + hook rejects each one) and re-raises, the request never reaches the + real LLM call -- neither async_log_success_event nor + async_log_failure_event, both tied to that call's own wrapper, ever + fires. async_post_call_failure_hook is the only remaining release path + for a reservation this hook already admitted earlier in that same + fallback chain. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + # This hook admits and reserves the only slot. + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-a"}, + call_type="completion", + ) + # _pre_call_with_fallbacks eventually gives up (every fallback rejected + # by some other hook) and reports the failure via post_call_failure_hook. + await hook.async_post_call_failure_hook( + request_data={"litellm_call_id": "call-1"}, + original_exception=ProxyRateLimitError( + detail={"error": "some_other_hooks_limit"}, headers={}, rate_limit_type=None + ), + user_api_key_dict=_key(), + ) + + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**_data(["end_user_id:u1"], call_id="call-2"), "model": "model-a"}, + call_type="completion", + ) + assert result is not None + + @pytest.mark.asyncio async def test_concurrency_reservation_released_when_disconnect_runs_in_a_forked_task(time_controller, monkeypatch): """ @@ -1279,11 +1327,18 @@ async def test_success_accounting_also_resolves_identity_from_the_policy_backed_ ) hook = _make_hook(time_controller) + # kwargs at async_log_success_event time is Logging.model_call_details: + # metadata/inherited_tags live nested under kwargs["litellm_params"], + # never at the top level -- see + # test_log_success_event_accounts_when_litellm_params_carries_a_null_litellm_metadata_key + # for the same shape. kwargs = { "litellm_call_id": "attack-1", - "metadata": { - "tags": ["company_id:attacker-chosen", "company_id:real-company"], - "inherited_tags": ["company_id:real-company"], + "litellm_params": { + "metadata": { + "tags": ["company_id:attacker-chosen", "company_id:real-company"], + "inherited_tags": ["company_id:real-company"], + }, }, "standard_logging_object": {"total_tokens": 0, "response_cost": 20.0}, } @@ -1302,3 +1357,46 @@ async def test_success_accounting_also_resolves_identity_from_the_policy_backed_ }, call_type="completion", ) + + +@pytest.mark.asyncio +async def test_rejection_detail_does_not_disclose_the_resolved_tag_value(time_controller, monkeypatch): + """ + tag_value can resolve from inherited_tags (server-assigned key/team/ + project metadata via order_tags_for_identity_resolution), so echoing it + back in the client-facing 429 detail would disclose that identity to + the caller who triggered the rejection. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "per-company", "tag_id": "company_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + key = _key() + + data = { + "litellm_call_id": "call-1", + "metadata": {"tags": ["company_id:secret-internal-name"], "inherited_tags": ["company_id:secret-internal-name"]}, + } + await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=data, call_type="completion") + + with pytest.raises(ProxyRateLimitError) as exc_info: + await hook.async_pre_call_hook( + user_api_key_dict=key, + cache=DualCache(), + data={ + "litellm_call_id": "call-2", + "metadata": { + "tags": ["company_id:secret-internal-name"], + "inherited_tags": ["company_id:secret-internal-name"], + }, + }, + call_type="completion", + ) + + assert "tag_value" not in exc_info.value.detail From f8b0705cb45b96935a948caa54034d38a75c0fe0 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 17:29:54 -0400 Subject: [PATCH 41/44] fix(rate-limiting): refresh the global hook's own concurrency key ttl on every admission veria-ai finding on the previous commit: TAG_RL_CHECK_AND_INCR_SCRIPT is shared with model_based_tag_rate_limits_hook, but only that hook's own Redis and in-memory call sites were updated to carry refresh_ttl through. global_tag_rate_limits_hook's own _check_and_increment_one still called the script with three args and never passed refresh_ttl to InMemoryCache, so a global concurrency bucket's ttl stayed fixed from its first admission and could expire mid-flight under sustained traffic, admitting past the cap. Also fixes ANN001/reportArgumentType drift the async_log_success_event and async_log_failure_event annotations surfaced: kwargs/response_obj are now typed, litellm_params_for_metadata is narrowed via isinstance instead of an unchecked assumption, and a stray tuple() wrap that never matched its list-typed parameter is dropped. --- .../hooks/global_tag_rate_limits_hook.py | 49 +++++--- .../hooks/test_global_tag_rate_limits_hook.py | 109 ++++++++++++++++++ 2 files changed, 143 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index f0a3c481ab6..09c2230338c 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -120,7 +120,6 @@ from litellm.router_strategy.tag_based_routing import ( ) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import TagRateLimitEntry, TagRateLimits -from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -327,10 +326,12 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o return built async def _check_and_increment_one( - self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int + self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool ) -> tuple[bool, float]: if self._check_and_incr_script is not None: - raw: Final = await self._check_and_incr_script(keys=(key,), args=(limit, increment, ttl)) + raw: Final = await self._check_and_incr_script( + keys=(key,), args=(limit, increment, ttl, 1 if refresh_ttl else 0) + ) return bool(raw[0]), float(raw[1]) async with self._lock: current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None) @@ -338,7 +339,9 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o 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) + await cache.async_set_cache( + key=key, value=new_value, ttl=ttl, refresh_ttl=refresh_ttl, litellm_parent_otel_span=None + ) return True, new_value async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None: @@ -352,17 +355,17 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o async def _atomic_check_and_increment( self, - checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], + checks: Sequence[tuple[InternalUsageCache, str, float, float, int, bool]], ) -> tuple[int | None, tuple[float, ...]]: """All-or-nothing atomic admission across `checks`: on a rejection, refunds every check admitted earlier in this batch.""" if not checks: return None, () admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection - for index, (cache, key, limit, increment, ttl) in enumerate(checks): + for index, (cache, key, limit, increment, ttl, refresh_ttl) in enumerate(checks): admitted = False try: - admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl) + admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl) finally: if not admitted: await self._refund_admitted(checks, up_to_index=index) @@ -373,10 +376,10 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o return None, tuple(admitted_values) async def _refund_admitted( - self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], up_to_index: int + self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int, bool]], up_to_index: int ) -> None: for refund_index in range(up_to_index): - refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index] + refund_cache, refund_key, _limit, refund_increment, _ttl, _refresh_ttl = checks[refund_index] try: await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment) except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest @@ -597,6 +600,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o ) else 1.0, self._ttl_for(check.unit, check.entry), + check.unit == "concurrency", ) for partition, check in zip(atomic_partitions, atomic_checks) ) @@ -667,7 +671,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o """ await self._release_pending_for_call_id(request_data) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime | None, + end_time: datetime | None, + ) -> None: # Always release regardless of which hook raised: this hook's own # rejection never reserves a slot, so pending_concurrency_keys is # already empty in that case and the check below no-ops; a rejection @@ -675,7 +685,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o # land after this hook already reserved its own slot. await self._release_pending_for_call_id(kwargs) - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime | None, + end_time: datetime | None, + ) -> None: stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) if stash is not None and stash.pending_concurrency_keys: release_keys: Final = tuple(stash.pending_concurrency_keys) @@ -688,15 +704,18 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o if config is None: return - standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - if standard_logging_object is None: + standard_logging_object: Final = kwargs.get("standard_logging_object") + if not isinstance(standard_logging_object, dict): return # kwargs here is Logging.model_call_details, not the router's flat # request kwargs admission sees: metadata/litellm_metadata are never # top-level here, only nested under kwargs["litellm_params"] (see # Logging.update_environment_variables). - litellm_params_for_metadata: Final = kwargs.get("litellm_params") or kwargs + litellm_params_raw: Final = kwargs.get("litellm_params") + litellm_params_for_metadata: Final[Mapping[str, object]] = ( + litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs + ) metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata) key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name) @@ -761,7 +780,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration accounting_task = asyncio.create_task( # not Final: rebound each loop iteration partition.v3.async_increment_tokens_with_ttl_preservation( - pipeline_operations=tuple(group_operations), parent_otel_span=None + pipeline_operations=group_operations, parent_otel_span=None ) ) _BACKGROUND_TASKS.add(accounting_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index c77cde23484..6f0d7d1aaf2 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -3,6 +3,9 @@ Unit tests for the global-scope, model-independent tag rate limiter. """ import asyncio +import os +import time +import uuid from datetime import datetime, timedelta import pytest @@ -48,6 +51,18 @@ def _data(tags: list[str], call_id: str = "call-1") -> dict: return {"metadata": {"tags": tags}, "litellm_call_id": call_id} +def _redis_hook(time_controller: TimeController): + from litellm.caching.redis_cache import RedisCache + + redis_host = os.getenv("REDIS_HOST") + redis_port = os.getenv("REDIS_PORT") + if not redis_host or not redis_port: + pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set") + redis_cache = RedisCache(host=redis_host, port=int(redis_port), password=os.getenv("REDIS_PASSWORD")) + dual_cache = DualCache(redis_cache=redis_cache) + return _PROXY_GlobalTagRateLimitsHook(internal_usage_cache=dual_cache, time_provider=time_controller.now), redis_cache + + # --------------------------------------------------------------------------- # No-op when unconfigured # --------------------------------------------------------------------------- @@ -1400,3 +1415,97 @@ async def test_rejection_detail_does_not_disclose_the_resolved_tag_value(time_co ) assert "tag_value" not in exc_info.value.detail + + +# --------------------------------------------------------------------------- +# Concurrency ttl refresh -- veria-ai finding: this hook's own atomic checks +# never carried refresh_ttl through to TAG_RL_CHECK_AND_INCR_SCRIPT or +# InMemoryCache.set_cache, even though it shares that script with +# model_based_tag_rate_limits_hook (whose own call sites got fixed first) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + A concurrency bucket isn't epoch-windowed like requests/tokens/dollars -- + its ttl exists only as a crash-safety net for a reservation whose + explicit release never runs -- so a still-active bucket receiving + continuous admissions must keep extending that ttl, or it expires + mid-flight under sustained traffic, silently admitting past the cap. + """ + hook, redis_cache = _redis_hook(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + key = f"{{tag_rl:test:global-ttl-refresh:{uuid.uuid4().hex}}}:inflight" + cache = hook.internal_usage_cache + try: + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = await redis_cache.init_async_client().ttl(key) + assert ttl_after_first_admission > 0 + + await asyncio.sleep(2) + + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = await redis_cache.init_async_client().ttl(key) + assert ttl_after_second_admission >= 2 + finally: + await redis_cache.async_delete_cache(key=key) + + +@pytest.mark.asyncio +async def test_in_memory_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + In-memory mirror of the Redis test above: InMemoryCache.allow_ttl_override + leaves a still-live ttl untouched, so without refresh_ttl reaching + set_cache a concurrency counter's expiry stayed fixed from its first + admission even under sustained traffic. + """ + hook = _make_hook(time_controller) + cache = hook.internal_usage_cache + in_memory_cache = cache.dual_cache.in_memory_cache + key = f"tag_rl:test:global-in-memory-ttl-refresh:{uuid.uuid4().hex}" + + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = in_memory_cache.ttl_dict[key] + + admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = in_memory_cache.ttl_dict[key] + + assert ttl_after_second_admission > ttl_after_first_admission + + +@pytest.mark.asyncio +async def test_concurrency_limit_admission_refreshes_ttl_end_to_end(time_controller, monkeypatch): + """ + End-to-end regression through async_pre_call_hook itself (not just the + low-level _check_and_increment_one helper above): a concurrency entry's + bucket key must carry a live ttl after admission, proving refresh_ttl is + actually wired from the classified check through to the atomic batch, + not just present on the helper's own signature. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 5, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion" + ) + + in_memory_cache = hook.internal_usage_cache.dual_cache.in_memory_cache + inflight_keys = [key for key in in_memory_cache.ttl_dict if key.endswith(":inflight")] + assert len(inflight_keys) == 1 + assert in_memory_cache.ttl_dict[inflight_keys[0]] > time.time() From ddbaa0daa3391206ec659e96b2325b28b7f917b7 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 17:38:57 -0400 Subject: [PATCH 42/44] fix(rate-limiting): close the same caller-forged metadata bypass at admission The base rename (resolve_success_event_metadata_variable_name -> resolve_authoritative_metadata_variable_name) broke this hook's import. Fixing the import also surfaced that this hook's own admission still used get_metadata_variable_name_from_kwargs, the naive key-presence check the sibling hook's own fix just replaced: a caller forging an empty (or None) litellm_metadata alongside the real, populated metadata made admission read no tags at all, admitting past every configured limit. Confirmed exploitable via a direct repro (5/5 requests admitted against a limit of 1) before the fix, blocked after. Regression test added. --- .../hooks/global_tag_rate_limits_hook.py | 12 ++++-- .../hooks/test_global_tag_rate_limits_hook.py | 39 +++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index 09c2230338c..f5d6857a18f 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -47,7 +47,6 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( @@ -112,7 +111,7 @@ 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, + resolve_authoritative_metadata_variable_name as _resolve_authoritative_metadata_variable_name, ) from litellm.proxy.utils import InternalUsageCache from litellm.router_strategy.tag_based_routing import ( @@ -544,7 +543,12 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o # re-charging. stash: Final = _claim_stash_for_data(data) - metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(data) + # Not get_metadata_variable_name_from_kwargs (naive key-presence + # check): a caller can forge an empty (or None) litellm_metadata on + # an ordinary request to make that check pick it over the real, + # populated metadata the proxy wrote identity/tags into, seeing no + # tags at all and admitting past every configured limit. + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(data) tags: Final = _order_tags_for_identity_resolution( _get_tags_from_request_kwargs(data, metadata_variable_name=metadata_variable_name), data, @@ -716,7 +720,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o litellm_params_for_metadata: Final[Mapping[str, object]] = ( 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) + metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name) diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 6f0d7d1aaf2..1a43bebb3b8 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -1509,3 +1509,42 @@ async def test_concurrency_limit_admission_refreshes_ttl_end_to_end(time_control inflight_keys = [key for key in in_memory_cache.ttl_dict if key.endswith(":inflight")] assert len(inflight_keys) == 1 assert in_memory_cache.ttl_dict[inflight_keys[0]] > time.time() + + +# --------------------------------------------------------------------------- +# Admission must resolve the authoritative metadata bucket, not the naive +# key-presence check -- veria-ai finding on the sibling model-based hook, +# same vulnerability class in this hook's own admission call site +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_admission_ignores_a_forged_empty_litellm_metadata_key(time_controller, monkeypatch): + """ + get_metadata_variable_name_from_kwargs picks "litellm_metadata" whenever + that key is merely present, regardless of its value. A caller adding an + empty "litellm_metadata" alongside the real, populated "metadata" made + admission read no tags at all, sailing past every configured limit. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "request_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + poisoned = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": {}} + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), cache=DualCache(), data={**poisoned, "litellm_call_id": "call-1"}, call_type="completion" + ) + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data={**poisoned, "litellm_call_id": "call-2"}, + call_type="completion", + ) From 2b0c1699d42bdd6240dde5598b534decea45302c Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 20:15:03 -0400 Subject: [PATCH 43/44] fix(rate-limiting): read success-event tags from the resolved metadata bucket, not raw kwargs Bugbot finding: async_log_success_event passed raw kwargs (Logging.model_call_details) to _get_tags_from_request_kwargs, which checks a top-level key matching metadata_variable_name before falling back to kwargs["litellm_params"]. Some call paths populate that top-level key present but None, so it won it over the real, populated dict nested under litellm_params, silently dropping token/dollar accounting for those routes. key_hash/key_alias/order_tags_for_identity_resolution already read from litellm_params_for_metadata; tag extraction now does too. Regression test confirmed to fail on the pre-fix code. --- .../hooks/global_tag_rate_limits_hook.py | 2 +- .../hooks/test_global_tag_rate_limits_hook.py | 44 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index f5d6857a18f..43d2a1462f2 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -725,7 +725,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name) tags: Final = _order_tags_for_identity_resolution( - _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name), + _get_tags_from_request_kwargs(litellm_params_for_metadata, metadata_variable_name=metadata_variable_name), litellm_params_for_metadata, metadata_variable_name, ) diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 1a43bebb3b8..aea37ecf149 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -1181,6 +1181,50 @@ async def test_log_success_event_accounts_when_litellm_params_carries_a_null_lit ) +@pytest.mark.asyncio +async def test_log_success_event_reads_tags_when_top_level_kwargs_carries_a_null_metadata_key( + time_controller, monkeypatch +): + """ + Bugbot finding: some call paths populate a top-level "metadata" (or + "litellm_metadata") key on kwargs (Logging.model_call_details) set to + None, alongside the real, populated dict nested under + kwargs["litellm_params"]. _get_tags_from_request_kwargs checks the + top-level key first; passing it raw kwargs made a present-but-None + top-level key win, reading no tags at all even though + metadata_variable_name correctly resolved to "metadata". + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "dollar_limits": { + "limits": [{"name": "daily_spend", "tag_id": "end_user_id", "limit": 10.0, "period_seconds": 86400}] + } + }, + ) + hook = _make_hook(time_controller) + + kwargs = { + "litellm_call_id": "call-1", + "metadata": None, + "litellm_params": { + "metadata": {"tags": ["end_user_id:u1"], "user_api_key": "hash"}, + }, + "standard_logging_object": {"total_tokens": 0, "response_cost": 12.0}, + } + await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + with pytest.raises(ProxyRateLimitError): + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + + @pytest.mark.asyncio async def test_dollar_limit_respects_apply_to_key_alias_at_accounting_time(time_controller, monkeypatch): """The entry only applies to `premium-key`; a non-listed key's spend must From 1c6e2f795919ca405ef2cba37899fbfb0afe7c1d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 20:34:00 -0400 Subject: [PATCH 44/44] fix(proxy): restore keepalive-ping exclusion from streaming disconnect refund An earlier round's rebase replayed a stale, pre-fix version of a commit this branch had already picked up the fix for, reverting delivered_chunk's keepalive-exclusion check and the _withheld_provider_output guard back to their original, buggy form. A client disconnecting after only keepalive pings (or while an agentic stream holds back real output) was refunded to input cost even though billable output had already been generated, exactly the veria-ai finding already fixed once upstream. Restores both checks; test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_cost passes again. --- litellm/proxy/common_request_processing.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a687aaf2c5c..f98ff8e8e62 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,6 +29,7 @@ from litellm.constants import ( NON_INFERENCE_CALL_TYPES, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, + STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -203,6 +204,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma } +def _withheld_provider_output(response: object) -> bool: + return getattr(response, "has_buffered_provider_output", False) is True + + def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -3497,8 +3502,9 @@ class ProxyBaseLLMRequestProcessing: # so a GeneratorExit on client disconnect is raised there and any # statement after the yield never runs. The slow-path hook is # awaited above, so a cancellation during it still leaves this - # False and refunds. - delivered_chunk = True + # False and refunds. A keepalive ping carries no provider output, + # so it must not suppress that refund. + delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) stream_completed = True except (asyncio.CancelledError, GeneratorExit): @@ -3512,7 +3518,7 @@ class ProxyBaseLLMRequestProcessing: # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: client_disconnected = True - if not delivered_chunk: + if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, )