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..c16003cd3bf --- /dev/null +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -0,0 +1,362 @@ +""" +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. + + A plain truthiness check on `litellm_metadata` isn't enough either: a + caller can populate it with unrelated, non-empty content on a route where + `metadata` is the field the proxy actually wrote identity into, and + truthiness alone would still misresolve to the caller-controlled bucket. + `add_user_api_key_auth_to_request_metadata` (litellm_pre_call_utils.py) + unconditionally stamps a `user_api_key_auth` key into whichever bucket it + resolved as authoritative, overwriting anything a caller pre-populated + there -- so requiring that marker's presence, not mere truthiness, only + ever prefers `litellm_metadata` when it is genuinely the field the proxy + wrote identity/tags into.""" + litellm_metadata: Final = litellm_params_for_metadata.get("litellm_metadata") + if isinstance(litellm_metadata, Mapping) and "user_api_key_auth" in litellm_metadata: + return "litellm_metadata" + return "metadata" + + +def extract_key_hash(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Same single-authoritative-field lookup as + `model_based_tag_rate_limits_hook._extract_team_id`, but for the calling + virtual key's hash: `LiteLLMProxyRequestSetup` sets + `metadata["user_api_key"]` to `user_api_key_dict.api_key`, which despite + the plain name is already the hashed token (see `litellm_pre_call_utils.py`). + """ + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + key_hash: Final = active.get("user_api_key") + return key_hash if isinstance(key_hash, str) else None + + +def extract_key_alias(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Same single-authoritative-field lookup as + `model_based_tag_rate_limits_hook._extract_team_id`, but for the calling + virtual key's own `key_alias`: `LiteLLMProxyRequestSetup` sets + `metadata["user_api_key_alias"]` to `user_api_key_dict.key_alias` + (see `litellm_pre_call_utils.py`).""" + active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING + key_alias: Final = active.get("user_api_key_alias") + return key_alias if isinstance(key_alias, str) else None + + +def 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 + a caller-supplied tag) with no length or content bound. Embedding it + directly would let a caller inflate a hook's own in-memory dict keys + past what `max_in_memory_cache_size` bounds (that caps item *count*, not + key bytes) and grow unbounded Redis keys with no cap at all. Hashing to + a fixed-length digest bounds a hook's own contribution to key size + regardless of the caller's input, while still preserving distinctness + (two different tag values still resolve to two different buckets). + """ + return hashlib.sha256(tag_value.encode()).hexdigest() + + +def policy_fingerprint(entry: TagRateLimitEntry) -> str: + """ + Two entries can share a `name` and `tag_id` while genuinely disagreeing + on `limit`, `period_seconds`, `scope_by_key_hash`, or any of the scoping + fields -- both hooks already treat that as two distinct policies + elsewhere (model_based_tag_rate_limits_hook.py's own per-deployment + dedup signature is one example), so the Redis/in-memory bucket key must + too, or two differently-configured entries that happen to share a name + check and charge the identical counter. `scope_by_key_hash` specifically + needs its own slot here rather than relying on a hook's own `_hash_tag` + `key_hash`-derived suffix to carry it: that suffix is empty whenever + `key_hash` resolves to `None` (no virtual key on the call), which would + otherwise collide an unscoped entry with a key-hash-scoped one that + agrees on every other field. Hashed to a fixed-length digest for the + same reason `fixed_length_identity` hashes `tag_value`: an operator's + own `enabled_for`/`disabled_for`/`apply_to_key_alias`/`apply_to_models` + list has no length bound. + """ + fingerprint_source: Final = ( + entry.limit, + entry.period_seconds, + entry.scope_by_key_hash, + scope_signature(entry.enabled_for), + scope_signature(entry.disabled_for), + entry.apply_to_key_alias, + entry.apply_to_models, + ) + return hashlib.sha256(repr(fingerprint_source).encode()).hexdigest()[:16] + + +def bucket_ttl_seconds(entry: TagRateLimitEntry) -> int: + """Redis (and in-memory fallback) TTL for a non-concurrency bucket key. + `entry.key_ttl_seconds` overrides the default of period_seconds + 3600 + when set -- see TagRateLimitEntry.key_ttl_seconds.""" + return entry.key_ttl_seconds if entry.key_ttl_seconds is not None else entry.period_seconds + 3600 + + +# None => this entry shares the hook's single default cache partition +# (matching every entry's behavior before this override existed). Otherwise +# a value-stable signature -- not the override int alone -- so two different +# entries that happen to choose the identical max_in_memory_cache_size don't +# get merged into one shared partition; the same entry (same config content) +# always resolves to the same signature across config/index rebuilds, which +# is what keeps each hook's own `_partitions` cache from leaking a fresh +# partition every time its configuration is re-resolved. +# The str before the trailing int is `policy_fingerprint(entry)`: two +# entries can share tag_id/name while genuinely disagreeing on +# limit/period_seconds/scope_by_key_hash/enabled_for/disabled_for/ +# apply_to_key_alias/apply_to_models -- both hooks already treat that as two +# distinct policies for bucket-key purposes, so a shared +# max_in_memory_cache_size must not route them onto the same partition +# either, or one entry's high-cardinality traffic can evict the other's +# active counters from a cache neither entry asked to share. +# max_in_memory_cache_size stays the trailing element: `partition_key[-1]` +# reads it directly to size the partition's cache. +PartitionKey: TypeAlias = tuple[str, str, str, int] | None +# Grouping type for async_log_success_event's per-partition tokens/dollars +# pipeline dispatch -- named only so the declaration fits on one line; see +# that method for why the grouping is needed. +PartitionOperations: TypeAlias = dict[PartitionKey, list[RedisPipelineIncrementOperation]] + + +def partition_key(entry: TagRateLimitEntry) -> PartitionKey: + if entry.max_in_memory_cache_size is None: + return None + return ( + entry.tag_id, + entry.name, + policy_fingerprint(entry), + entry.max_in_memory_cache_size, + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index a3335be2b2b..47eae2bf4a8 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum +import math from collections.abc import Mapping from dataclasses import dataclass from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -137,6 +138,176 @@ def _as_utc(value: datetime.datetime | None) -> datetime.datetime | None: return value.astimezone(datetime.timezone.utc) +class TagRateLimitScope(BaseModel): + """ + A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + entry to `tag_id: company_id, values: ["1032"]` so it only applies to + requests tagged as belonging to company 1032, independent of whichever + tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + `disabled_for`, which are the only two fields that construct this. + """ + + tag_id: str + values: tuple[str, ...] + + model_config = ConfigDict(frozen=True) + + @model_validator(mode="after") + def _validate_values(self) -> "TagRateLimitScope": + if not self.values: + raise ValueError("values must be a non-empty list of strings") + return self + + @model_validator(mode="after") + def _normalize_values(self) -> "TagRateLimitScope": + # Sorted and deduplicated: `values` is only ever used for membership + # tests (see _entry_applies), never order-dependent, but is also + # folded verbatim into the dedup signature two deployments' entries + # are compared by (see _scope_signature) -- an unsorted tuple would + # make config-order alone, not policy, decide whether two entries + # dedup to one shared bucket or wrongly split into two. + # object.__setattr__ bypasses this frozen model's own assignment + # guard -- returning a replacement instance from an "after" validator + # is silently ignored when constructing via __init__ (only takes + # effect via model_validate), so mutating in place is the only way + # this normalization reliably applies regardless of construction path. + object.__setattr__(self, "values", tuple(sorted(set(self.values)))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitEntry(BaseModel): + name: str + tag_id: str = "end_user_id" + limit: float + period_seconds: int + scope_by_key_hash: bool = False + # Overrides this entry's bucket/reservation key TTL (Redis, and the + # in-memory fallback when Redis isn't configured). Defaults to + # period_seconds + 3600 when unset -- see _PROXY_ModelBasedTagRateLimitsHook._ttl_for. + # A high-cardinality tag_id can keep many keys alive at once; lowering + # this lets an operator shed them sooner without shortening + # period_seconds itself. + key_ttl_seconds: int | None = None + # Overrides the size of the dedicated in-memory cache partition this + # entry's own keys live in, when Redis isn't configured (or as a local + # fast-path cache when it is). Unset means this entry shares the hook's + # single default partition, sized by + # litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200 if that's also + # unset). A high-cardinality tag_id can churn past that shared cap and + # evict another entry's active counters; setting this gives the entry + # its own dedicated partition instead. + max_in_memory_cache_size: int | None = None + # Gate this entry on a tag -- often a SECOND, independent tag (e.g. + # `enabled_for: {tag_id: company_id, values: ["1032"]}` to scope an + # override to one company's traffic), but `tag_id` can equally be set to + # this same entry's own `tag_id` to scope by a subset of its own + # resolved identity instead, without a second tag at all. + # `disabled_for` is checked first (deny overrides allow) when both are + # set. An absent gate tag never satisfies `enabled_for` (an allowlist + # gate requires an explicit match) but never triggers `disabled_for` + # either (nothing to match against a denylist). + enabled_for: TagRateLimitScope | None = None + disabled_for: TagRateLimitScope | None = None + # Restrict this entry to requests authenticated with one of these virtual + # keys' own `key_alias`. Unset (the default) means the entry applies to + # every request regardless of which key made it. A key with no alias set + # never satisfies this allowlist, same "absent gate never matches an + # allowlist" precedent as `enabled_for`. + apply_to_key_alias: tuple[str, ...] | None = None + # Restrict this entry to requests whose caller-facing `model` matches one + # of these names. Unset (the default) means the entry applies to every + # model. A request with no `model` field never satisfies this allowlist, + # same "absent gate never matches an allowlist" precedent as + # `apply_to_key_alias`. + apply_to_models: tuple[str, ...] | None = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _validate_limit(self) -> "TagRateLimitEntry": + # NaN compares False against every ordering operator, so a NaN limit + # makes the atomic requests/concurrency check-and-increment (which + # rejects when the new value exceeds the limit) never reject -- + # admitting indefinitely -- while the read-only tokens/dollars check + # (which admits when the current value is under the limit) never + # admits, rejecting every tagged request. Either outcome silently + # defeats the entry; reject it at config load time instead. + if math.isnan(self.limit): + raise ValueError("limit must not be NaN") + # Positive infinity never rejects the checks that gate this limit; negative + # infinity always does. Both silently defeat the entry. + if math.isinf(self.limit): + raise ValueError("limit must be finite") + # Zero or negative makes every check that gates this limit either always + # reject or never admit, silently blocking or admitting all matching traffic + # instead of the likely intended config. + if self.limit <= 0: + raise ValueError("limit must be a positive number") + return self + + @model_validator(mode="after") + def _validate_period_seconds(self) -> "TagRateLimitEntry": + if self.period_seconds <= 0: + raise ValueError("period_seconds must be a positive integer") + return self + + @model_validator(mode="after") + def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": + if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: + raise ValueError("key_ttl_seconds must be a positive integer when set") + # A shorter TTL than period_seconds expires the counter before its period + # elapses, letting it reset early and exceed the limit. + if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: + raise ValueError("key_ttl_seconds must be at least period_seconds when set") + return self + + @model_validator(mode="after") + def _validate_max_in_memory_cache_size(self) -> "TagRateLimitEntry": + if self.max_in_memory_cache_size is not None and self.max_in_memory_cache_size <= 0: + raise ValueError("max_in_memory_cache_size must be a positive integer when set") + return self + + @model_validator(mode="after") + def _validate_apply_to_key_alias(self) -> "TagRateLimitEntry": + if self.apply_to_key_alias is not None and not self.apply_to_key_alias: + raise ValueError("apply_to_key_alias must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_key_alias(self) -> "TagRateLimitEntry": + # Sorted and deduplicated for the same reason as + # TagRateLimitScope._normalize_values: only ever used for membership + # tests, but also folded verbatim into the dedup signature, where an + # unsorted tuple would make config-order alone decide whether two + # deployments' entries dedup to one shared bucket. + if self.apply_to_key_alias is not None: + self.apply_to_key_alias = tuple(sorted(set(self.apply_to_key_alias))) # mutable-ok: frozen before escaping + return self + + @model_validator(mode="after") + def _validate_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None and not self.apply_to_models: + raise ValueError("apply_to_models must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None: + self.apply_to_models = tuple(sorted(set(self.apply_to_models))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitGroup(BaseModel): + limits: tuple[TagRateLimitEntry, ...] = () + + +class TagRateLimits(BaseModel): + token_limits: TagRateLimitGroup | None = None + request_limits: TagRateLimitGroup | None = None + dollar_limits: TagRateLimitGroup | None = None + concurrency_limits: TagRateLimitGroup | None = None + + class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. @@ -186,6 +357,8 @@ class ModelInfo(MirroredPricingParams): # router-wide default. enable_tag_filtering: bool | None = None + tag_rate_limits: TagRateLimits | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/tests/test_litellm/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..7bb9d04fef0 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py @@ -0,0 +1,421 @@ +""" +Unit tests for the primitives shared by both tag-scoped rate-limit hooks +(`model_based_tag_rate_limits_hook.py` and `global_tag_rate_limits_hook.py`). + +These test the hook-independent logic in isolation: identity extraction, +`entry_applies` scoping, and the partition/bucket-TTL key helpers. Each +hook's own test file covers everything specific to how it wires these +primitives into its own admission/accounting engine. +""" + +from litellm.proxy.hooks.tag_rate_limits_shared import ( + bucket_ttl_seconds, + entry_applies, + extract_identity, + extract_key_hash, + fixed_length_identity, + order_tags_for_identity_resolution, + partition_key, + resolve_success_event_metadata_variable_name, +) +from litellm.types.router import TagRateLimitEntry, TagRateLimitScope + +# --------------------------------------------------------------------------- +# extract_identity +# --------------------------------------------------------------------------- + + +def test_extract_identity_matches_prefixed_tag(): + assert extract_identity(["team_id:t1", "end_user_id:u1"], "end_user_id") == "u1" + + +def test_extract_identity_returns_none_when_absent(): + assert extract_identity(["team_id:t1"], "end_user_id") is None + + +def test_extract_identity_skips_negation_tags(): + """A `!end_user_id:u1` routing-negation marker must never be read as identity.""" + assert extract_identity(["!end_user_id:u1"], "end_user_id") is None + + +# --------------------------------------------------------------------------- +# order_tags_for_identity_resolution -- veria-ai finding on PR #38292: a +# caller-supplied tag must not shadow a policy-backed (key/team/project) +# tag sharing the same tag_id prefix +# --------------------------------------------------------------------------- + + +def test_order_tags_for_identity_resolution_prefers_inherited_tag_over_caller_supplied(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["company_id:attacker-chosen", "end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "company_id") == "real-company" + + +def test_order_tags_for_identity_resolution_falls_back_to_caller_tags_when_nothing_inherited(): + request_kwargs = {"metadata": {}} + tags = ["end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "end_user_id") == "u1" + + +def test_order_tags_for_identity_resolution_keeps_caller_only_tags_not_shadowed_by_a_different_tag_id(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["end_user_id:u1"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert extract_identity(ordered, "end_user_id") == "u1" + assert extract_identity(ordered, "company_id") == "real-company" + + +def test_order_tags_for_identity_resolution_deduplicates_identical_tag_present_in_both_sources(): + request_kwargs = {"metadata": {"inherited_tags": ["company_id:real-company"]}} + tags = ["company_id:real-company"] + ordered = order_tags_for_identity_resolution(tags, request_kwargs, "metadata") + assert ordered.count("company_id:real-company") == 1 + + +# --------------------------------------------------------------------------- +# fixed_length_identity -- tag_value is caller-controlled with no length +# bound; a hook's own contribution to a cache key must not grow with it +# --------------------------------------------------------------------------- + + +def test_fixed_length_identity_bounds_key_contribution_regardless_of_input_size(): + """ + A caller can submit an arbitrarily long tag value (no length or content + bound is enforced upstream of either hook). Without hashing, that value + would go straight into an in-memory dict key (bypassing + max_in_memory_cache_size, which caps item *count* not key bytes) and an + unbounded-length Redis key (Redis has no key-count or key-size cap at + all here). A fixed-length digest bounds a hook's own contribution to + the key regardless of input size. + """ + huge_value = "x" * 5_000_000 + digest = fixed_length_identity(huge_value) + assert len(digest) == 64 # sha256 hex digest length, independent of input size + + +def test_fixed_length_identity_preserves_distinctness(): + """Hashing must not collapse two different tag values onto one bucket.""" + assert fixed_length_identity("user-a") != fixed_length_identity("user-b") + assert fixed_length_identity("user-a") == fixed_length_identity("user-a") + + +# --------------------------------------------------------------------------- +# extract_key_hash -- must read only the one field the server actually +# authenticates into, never fall back to the other +# --------------------------------------------------------------------------- + + +def test_extract_key_hash_ignores_a_forged_value_in_the_non_authoritative_field(): + """ + On a route where litellm_metadata is authoritative, the server writes + the real hash there and never touches metadata -- so a caller-supplied + metadata.user_api_key must not be read at all, let alone win. + """ + request_kwargs = { + "metadata": {"user_api_key": "forged-by-caller"}, + "litellm_metadata": {"user_api_key": "real-authenticated-hash"}, + } + assert extract_key_hash(request_kwargs, "litellm_metadata") == "real-authenticated-hash" + + +def test_extract_key_hash_reads_metadata_when_it_is_the_authoritative_field(): + request_kwargs = {"metadata": {"user_api_key": "real-hash"}} + assert extract_key_hash(request_kwargs, "metadata") == "real-hash" + + +# --------------------------------------------------------------------------- +# resolve_success_event_metadata_variable_name -- veria-ai finding surfaced on +# PR #38347: a caller-supplied, non-empty litellm_metadata must not be +# selected over the metadata the proxy actually wrote authenticated tags into +# --------------------------------------------------------------------------- + + +def test_resolve_success_event_metadata_variable_name_ignores_caller_forged_non_empty_litellm_metadata(): + """A caller-supplied litellm_metadata with unrelated content (no + user_api_key_auth marker) must not be picked over metadata, even though + it is a non-empty dict.""" + litellm_params_for_metadata = {"litellm_metadata": {"marker": True}} + assert resolve_success_event_metadata_variable_name(litellm_params_for_metadata) == "metadata" + + +def test_resolve_success_event_metadata_variable_name_selects_litellm_metadata_when_server_written(): + """A genuinely server-populated litellm_metadata (LITELLM_METADATA_ROUTES) + always carries the user_api_key_auth marker stamped by + add_user_api_key_auth_to_request_metadata.""" + litellm_params_for_metadata = { + "litellm_metadata": {"user_api_key_auth": object(), "tags": ["team_id:t1"]} + } + assert resolve_success_event_metadata_variable_name(litellm_params_for_metadata) == "litellm_metadata" + + +def test_resolve_success_event_metadata_variable_name_defaults_to_metadata_when_litellm_metadata_absent(): + assert resolve_success_event_metadata_variable_name({}) == "metadata" + + +def test_resolve_success_event_metadata_variable_name_defaults_to_metadata_when_litellm_metadata_none(): + assert resolve_success_event_metadata_variable_name({"litellm_metadata": None}) == "metadata" + + +def test_resolve_success_event_metadata_variable_name_defaults_to_metadata_when_litellm_metadata_empty(): + assert resolve_success_event_metadata_variable_name({"litellm_metadata": {}}) == "metadata" + + +# --------------------------------------------------------------------------- +# entry_applies -- enabled_for / disabled_for / apply_to_key_alias / apply_to_models +# --------------------------------------------------------------------------- + + +def test_entry_applies_with_none_of_the_scoping_fields_set(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_disabled_for_on_its_own_tag_id_excludes_a_listed_value(): + """disabled_for's `tag_id` can be set to the entry's own tag_id, gating on + a subset of its own resolved identity rather than a second tag.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + assert entry_applies(entry, ["end_user_id:u2"], None, None) is True + + +def test_entry_applies_enabled_for_on_its_own_tag_id_restricts_to_a_listed_value(): + """enabled_for's `tag_id` can likewise be set to the entry's own tag_id, + admitting only a hand-picked subset of its own resolved identity.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u2", "u3")), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + assert entry_applies(entry, ["end_user_id:u2"], None, None) is True + + +def test_entry_applies_matches_an_enabled_for_gate(): + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is True + + +def test_entry_applies_skips_when_enabled_for_gate_tag_is_absent(): + """ + enabled_for is an allowlist gate: absence of the gate tag must not + satisfy it, unlike disabled_for below. + """ + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_skips_when_disabled_for_gate_matches(): + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False + + +def test_entry_applies_when_disabled_for_gate_tag_is_absent(): + """disabled_for is a denylist gate: absence of the gate tag has nothing + to match against, so the entry still applies.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_disabled_for_overrides_a_matching_enabled_for_gate(): + """Deny (disabled_for) takes effect independently of whether the + enabled_for gate itself matched, even when both target the same tag.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)), + ) + assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False + + +def test_entry_applies_with_apply_to_key_alias_unset_applies_to_every_key(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], "any-key-alias", None) is True + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_admits_a_key_alias_on_the_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], "team-a-key", None) is True + + +def test_entry_applies_rejects_a_key_alias_missing_from_the_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], "team-b-key", None) is False + + +def test_entry_applies_rejects_when_key_has_no_alias_but_allowlist_is_set(): + """apply_to_key_alias is an allowlist gate: a key with no alias at all + never satisfies it, same as enabled_for's absent-gate-tag semantics.""" + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_with_apply_to_models_unset_applies_to_every_model(): + entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) + assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True + assert entry_applies(entry, ["end_user_id:u1"], None, None) is True + + +def test_entry_applies_admits_a_model_on_the_apply_to_models_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True + + +def test_entry_applies_rejects_a_model_missing_from_the_apply_to_models_allowlist(): + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, "sonnet-chain") is False + + +def test_entry_applies_rejects_when_model_is_absent_but_apply_to_models_is_set(): + """apply_to_models is an allowlist gate: a request with no model at all + never satisfies it, same as apply_to_key_alias's absent-key semantics.""" + entry = TagRateLimitEntry( + name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",) + ) + assert entry_applies(entry, ["end_user_id:u1"], None, None) is False + + +def test_entry_applies_apply_to_models_composes_with_apply_to_key_alias(): + """Both gates must pass: a request against the listed model but a + non-listed key alias must not apply, even though apply_to_models alone + would have admitted it.""" + entry = TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + apply_to_models=("opus-chain",), + apply_to_key_alias=("premium-key",), + ) + assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "opus-chain") is True + assert entry_applies(entry, ["end_user_id:u1"], "other-key", "opus-chain") is False + assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "sonnet-chain") is False + + +# --------------------------------------------------------------------------- +# partition_key -- entries that share max_in_memory_cache_size but disagree +# on any policy-fingerprinted field must never share a cache partition +# --------------------------------------------------------------------------- + + +def test_partition_key_distinguishes_entries_that_differ_only_by_scope_by_key_hash(): + """ + scope_by_key_hash is part of the partition-key signature: two entries + identical in every other field but differing only on this flag are + different rate limits (different bucket keys per each hook's own + `_hash_tag`) and must never be routed to the same cache partition. + """ + unscoped = TagRateLimitEntry( + name="per_minute", tag_id="end_user_id", limit=5, period_seconds=60, max_in_memory_cache_size=100 + ) + scoped = TagRateLimitEntry( + name="per_minute", + tag_id="end_user_id", + limit=5, + period_seconds=60, + scope_by_key_hash=True, + max_in_memory_cache_size=100, + ) + assert partition_key(unscoped) != partition_key(scoped) + + +def test_partition_key_distinguishes_entries_that_differ_only_by_scoping_fields(): + """ + A plain, unscoped entry and a scoped override can legitimately share + name/tag_id/limit/period_seconds/scope_by_key_hash while disagreeing on + enabled_for/disabled_for/apply_to_key_alias/apply_to_models -- + policy_fingerprint already treats that as two distinct policies, so a + shared max_in_memory_cache_size must not route them onto the same + in-memory partition either, or one entry's high-cardinality traffic can + evict the other's active counters from a cache neither entry asked to + share. + """ + base_kwargs = { + "name": "daily", + "tag_id": "end_user_id", + "limit": 100, + "period_seconds": 86400, + "max_in_memory_cache_size": 50, + } + unscoped = TagRateLimitEntry(**base_kwargs) + enabled_for_scoped = TagRateLimitEntry( + **base_kwargs, enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + disabled_for_scoped = TagRateLimitEntry( + **base_kwargs, disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + alias_scoped = TagRateLimitEntry(**base_kwargs, apply_to_key_alias=("premium-key",)) + models_scoped = TagRateLimitEntry(**base_kwargs, apply_to_models=("opus-chain",)) + + keys = { + partition_key(unscoped), + partition_key(enabled_for_scoped), + partition_key(disabled_for_scoped), + partition_key(alias_scoped), + partition_key(models_scoped), + } + assert len(keys) == 5 + + +# --------------------------------------------------------------------------- +# bucket_ttl_seconds -- per-tag Redis/bucket key TTL override +# --------------------------------------------------------------------------- + + +def test_bucket_ttl_seconds_defaults_to_period_plus_one_hour_when_unset(): + entry = TagRateLimitEntry(name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60) + assert bucket_ttl_seconds(entry) == 60 + 3600 + + +def test_bucket_ttl_seconds_honors_key_ttl_seconds_override(): + entry = TagRateLimitEntry( + name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60, key_ttl_seconds=120 + ) + assert bucket_ttl_seconds(entry) == 120 diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index accd3b32a0d..f58bd9b118e 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -5,6 +5,7 @@ from litellm.types.router import ( Deployment, LiteLLM_Params, ModelInfo, + TagRateLimitEntry, ) from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams @@ -89,3 +90,47 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") + + +# litellm/types/router.py is imported by plain SDK users, not just the proxy, so a +# TagRateLimitEntry validation error should read as a generic config-validation +# message and not describe the proxy rate-limit hook's internal admission mechanics. +_INTERNAL_ENFORCEMENT_JARGON = ( + "admission", + "tagged request", + "tagged traffic", + "check-and-increment", + "read-only", + "atomic", + "window rolls over", +) + + +def _assert_message_has_no_internal_jargon(excinfo: pytest.ExceptionInfo) -> None: + message = str(excinfo.value).lower() + leaked = [term for term in _INTERNAL_ENFORCEMENT_JARGON if term in message] + assert not leaked, f"validation message leaked internal enforcement jargon: {leaked}" + + +def test_limit_infinite_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=float("inf"), period_seconds=60) + _assert_message_has_no_internal_jargon(excinfo) + + +def test_limit_non_positive_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=0, period_seconds=60) + _assert_message_has_no_internal_jargon(excinfo) + + +def test_key_ttl_seconds_shorter_than_period_rejected_without_internal_enforcement_jargon(): + with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo: + TagRateLimitEntry( + name="daily", + tag_id="end_user_id", + limit=500, + period_seconds=86400, + key_ttl_seconds=60, + ) + _assert_message_has_no_internal_jargon(excinfo) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 405ec9a01bf..973adce672e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35115,6 +35115,64 @@ export interface components { /** Tpm Limit */ tpm_limit?: number | null; }; + /** TagRateLimitEntry */ + TagRateLimitEntry: { + /** Apply To Key Alias */ + apply_to_key_alias?: string[] | null; + /** Apply To Models */ + apply_to_models?: string[] | null; + disabled_for?: components["schemas"]["TagRateLimitScope"] | null; + enabled_for?: components["schemas"]["TagRateLimitScope"] | null; + /** Key Ttl Seconds */ + key_ttl_seconds?: number | null; + /** Limit */ + limit: number; + /** Max In Memory Cache Size */ + max_in_memory_cache_size?: number | null; + /** Name */ + name: string; + /** Period Seconds */ + period_seconds: number; + /** + * Scope By Key Hash + * @default false + */ + scope_by_key_hash: boolean; + /** + * Tag Id + * @default end_user_id + */ + tag_id: string; + }; + /** TagRateLimitGroup */ + TagRateLimitGroup: { + /** + * Limits + * @default [] + */ + limits: components["schemas"]["TagRateLimitEntry"][]; + }; + /** + * TagRateLimitScope + * @description A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + * entry to `tag_id: company_id, values: ["1032"]` so it only applies to + * requests tagged as belonging to company 1032, independent of whichever + * tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + * `disabled_for`, which are the only two fields that construct this. + */ + TagRateLimitScope: { + /** Tag Id */ + tag_id: string; + /** Values */ + values: string[]; + }; + /** TagRateLimits */ + TagRateLimits: { + concurrency_limits?: components["schemas"]["TagRateLimitGroup"] | null; + dollar_limits?: components["schemas"]["TagRateLimitGroup"] | null; + request_limits?: components["schemas"]["TagRateLimitGroup"] | null; + token_limits?: components["schemas"]["TagRateLimitGroup"] | null; + }; /** * TagSummaryMetrics * @description Summary metrics for a tag @@ -37890,6 +37948,7 @@ export interface components { ptu_effective_from?: string | null; /** Ptu Effective To */ ptu_effective_to?: string | null; + tag_rate_limits?: components["schemas"]["TagRateLimits"] | null; /** Team Id */ team_id?: string | null; /** Team Public Model Name */