scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge e8e6ae682e into eb0e3f8c18

Browse source
This commit is contained in:
Deepanshu Lulla 2026-08-28 00:23:08 +00:00 • committed by GitHub
commit f3d6bf5b89
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1060 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

362
litellm/proxy/hooks/tag_rate_limits_shared.py Normal file
View file

@ -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,
)

173
litellm/types/router.py
View file

@ -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

421
tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py Normal file
View file

@ -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

45
tests/test_litellm/types/test_router.py
View file

@ -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)

59
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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 */