scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(rate-limiting): extract shared tag-rate-limit helpers into their own module

Browse source 
Both tag-scoped rate limiting hooks need the same identity/scope
extraction, policy fingerprinting, bucket-key hashing, and cache
partitioning primitives. Moving them into their own module lets a
model-independent global hook consume them without reaching into a
model-based hook's private internals, which is how the two hooks
previously shared this logic.
This commit is contained in:
Deepanshu 2026-08-25 21:53:52 -04:00
parent 4492799fad
commit a5bee80768
2 changed files with 680 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -0,0 +1,334 @@
"""
Primitives shared by both tag-scoped rate-limit hooks:
`model_based_tag_rate_limits_hook.py` (per-deployment limits nested under
`model_info.tag_rate_limits`, admitted once per routing hop) and
`global_tag_rate_limits_hook.py` (model-independent limits declared once
under `litellm_settings.global_tag_rate_limits`, admitted once per request).
Both hooks enforce the identical `TagRateLimitEntry` shape and need the same
identity/scope extraction, policy fingerprinting, bucket-key hashing, and
cache-partitioning primitives, so those live here rather than in either
hook's own module -- keeping neither hook reaching into the other's private
internals to reuse them. Every name here is this module's own public
interface (no leading underscore): each hook imports what it needs aliased
back to its own historical, underscore-prefixed local name (e.g.
`entry_applies as _entry_applies`), so this is a genuine export, not a
private symbol either hook reaches across a module boundary to grab.
"""
import asyncio
import hashlib
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from litellm.exceptions import RateLimitType
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.router import TagRateLimitEntry, TagRateLimitScope
LimitUnit: TypeAlias = Literal["tokens", "requests", "dollars", "concurrency"]
LIMIT_UNITS: Final[tuple[LimitUnit, ...]] = ("tokens", "requests", "dollars", "concurrency")
# Units whose admission must be atomic (check-and-increment in one Redis
# round trip) because the increment amount is known upfront (always 1).
# tokens/dollars can't be: real usage is only known after the response, so
# they stay a read-then-account-on-success check with a documented,
# unavoidable admit-vs-account race.
ATOMIC_UNITS: Final[frozenset[LimitUnit]] = frozenset({"requests", "concurrency"})
UNIT_TO_GROUP_FIELD: Final[Mapping[LimitUnit, str]] = MappingProxyType(
{
"tokens": "token_limits",
"requests": "request_limits",
"dollars": "dollar_limits",
"concurrency": "concurrency_limits",
}
)
UNIT_TO_RATE_LIMIT_TYPE: Final[Mapping[LimitUnit, RateLimitType]] = MappingProxyType(
{
"tokens": RateLimitType.TOKENS,
"requests": RateLimitType.REQUESTS,
"dollars": RateLimitType.BUDGET,
"concurrency": RateLimitType.CONCURRENT_REQUESTS,
}
)
# Shared read-only fallback for an absent/None mapping (request_kwargs,
# metadata, model_info, ...): avoids constructing a fresh mutable `{}` at
# every one of these call sites just to immediately call `.get()` on it.
EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
# `asyncio.create_task`'s own docs: "Save a reference to the result of this
# function, to avoid a task disappearing mid-execution. The event loop only
# keeps weak references to tasks. A task that isn't referenced elsewhere may
# get garbage collected at any time, even before it's done." The success path
# deliberately fires-and-forgets its concurrency release and its token/dollar
# accounting increment (unlike the failure/disconnect paths, which await
# concurrency release directly) to keep the hot success-response path from
# waiting on a Redis round trip; by the time either background task would
# run, the state it needs (popped pending keys, or the request's own usage
# figures) is only available in that task's own closure, so a collected
# task's work is unrecoverable, not just delayed. Holding a strong reference
# here until each task's own completion callback discards it is the standard
# fix, shared by every fire-and-forget task either hook creates.
BACKGROUND_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: see comment above
# Floor for a concurrency reservation's self-heal TTL, regardless of the
# configured period_seconds. A reservation that expires while its request is
# still genuinely in flight silently admits requests past the limit; this
# generous floor keeps that window far larger than any realistic request
# duration, at the cost of a leaked (crashed-worker) slot self-healing more
# slowly. period_seconds can still raise the TTL further, never lower it.
CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600
# Single-key atomic check-and-increment. Deliberately one key per script call
# (never a batch of differently-hash-tagged keys in one call): every tag_rl
# key carries its own self-contained {..} hash tag so unrelated buckets never
# forcibly co-locate on the same Redis Cluster shard, which means a single Lua
# invocation can never span more than one key's slot without risking a
# cross-slot error. All-or-nothing across a hop's multiple atomic checks
# (e.g. requests + concurrency checked together) is achieved in Python by
# calling this once per key and refunding every earlier admission in the same
# batch if a later one is rejected -- the same refund-on-rollback shape as
# `atomic_check_and_increment_by_n` in parallel_request_limiter_v3.py, applied
# per-key instead of per-descriptor since each key already is one hash-tag
# group by construction.
TAG_RL_CHECK_AND_INCR_SCRIPT: Final = """
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local increment = tonumber(ARGV[2])
local ttl = tonumber(ARGV[3])
local current = tonumber(redis.call('GET', key) or 0)
if current + increment > limit then
return { 0, current }
end
local new_value = redis.call('INCRBY', key, increment)
local current_ttl = redis.call('TTL', key)
if current_ttl == -1 and ttl > 0 then
redis.call('EXPIRE', key, ttl)
end
return { 1, new_value }
"""
# Atomic decrement that never leaves a counter negative. Used both to refund
# an earlier admission when a later key in the same batch is rejected, and to
# release a concurrency reservation -- floors at 0 so a decrement that can't
# be attributed to the exact reservation that caused it (see
# `_release_keys`'s docstring) degrades to under-counting rather than a
# negative counter that would admit unlimited requests. Floors via DEL, not
# `SET key 0`: releasing a reservation whose key already expired makes
# INCRBY recreate it with no TTL, and a plain SET would leave that recreated
# key permanently in Redis (SET clears any TTL); DEL removes it outright,
# which reads back identically to 0 everywhere this key is read (`GET key or 0`).
TAG_RL_DECR_FLOOR_ZERO_SCRIPT: Final = """
local key = KEYS[1]
local delta = tonumber(ARGV[1])
local new_value = redis.call('INCRBY', key, delta)
if new_value < 0 then
redis.call('DEL', key)
new_value = 0
end
return new_value
"""
# A (tag_id, values) pair mirroring TagRateLimitScope's own fields, used to
# fold `enabled_for`/`disabled_for` into a hashable form (a policy
# fingerprint, or a per-deployment dedup signature) without depending on
# TagRateLimitScope's own hashability.
ScopeSignature: TypeAlias = tuple[str, tuple[str, ...]] | None
def scope_signature(scope: TagRateLimitScope | None) -> ScopeSignature:
"""Normalizes a `TagRateLimitScope` into a plain, hashable tuple, so
`enabled_for`/`disabled_for` can be folded into a policy fingerprint (or
a per-deployment dedup signature) -- two entries disagreeing on either
field must be treated as genuinely different policies rather than
merged into one bucket."""
return None if scope is None else (scope.tag_id, scope.values)
def extract_identity(tags: Sequence[str], tag_id: str) -> str | None:
"""
First tag matching `f"{tag_id}:"`, value after the colon. Tags starting
with `!` are tag-routing negation markers, not identity tags, and are
skipped so they can never be misread as an identity value.
"""
prefix: Final = f"{tag_id}:"
for tag in tags:
if tag.startswith("!"):
continue
if tag.startswith(prefix):
return tag[len(prefix) :]
return None
def entry_applies(entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, model: str | None) -> bool:
"""
Applies `entry`'s own scoping fields (`enabled_for`/`disabled_for`/
`apply_to_key_alias`/`apply_to_models`), evaluated in this order -- deny
overrides allow, checked before any allowlist:
1. `disabled_for`: the gate tag (often a SECOND, independent tag, but
`disabled_for.tag_id` can equally be set to this entry's own
`tag_id` to gate on a subset of its own resolved identity) is
present and its value is in `disabled_for.values` -> doesn't apply.
Absent gate tag never triggers this -- nothing to match against a
denylist.
2. `enabled_for`: the gate tag is absent, or present but its value is
NOT in `enabled_for.values` -> doesn't apply. Unlike `disabled_for`,
absence DOES fail this check -- an allowlist gate requires an
explicit match, so "not tagged at all" means "not in scope".
3. `apply_to_models`: `model` is absent, or present but not in the
list -> doesn't apply. Same allowlist semantics as `enabled_for` --
a request with no `model` never satisfies this gate.
4. `apply_to_key_alias`: the calling key's own alias is absent, or
present but not in the list -> doesn't apply. Same allowlist
semantics as `enabled_for` -- a key with no alias set never
satisfies this gate.
An entry with none of these fields set always applies -- this is the
unscoped behavior every existing entry has today, unchanged.
"""
if entry.disabled_for is not None:
disabled_gate_value: Final = extract_identity(tags, entry.disabled_for.tag_id)
if disabled_gate_value is not None and disabled_gate_value in entry.disabled_for.values:
return False
if entry.enabled_for is not None:
enabled_gate_value: Final = extract_identity(tags, entry.enabled_for.tag_id)
if enabled_gate_value is None or enabled_gate_value not in entry.enabled_for.values:
return False
if entry.apply_to_models is not None and model not in entry.apply_to_models:
return False
if entry.apply_to_key_alias is None:
return True
return key_alias in entry.apply_to_key_alias
def resolve_success_event_metadata_variable_name(
litellm_params_for_metadata: Mapping[str, object],
) -> Literal["metadata", "litellm_metadata"]:
"""`get_metadata_variable_name_from_kwargs` only checks key presence, which
misresolves at `async_log_success_event` time: `kwargs["litellm_params"]`
always carries a `litellm_metadata` key (typically `None`) alongside the
real, populated `metadata` dict for a standard (non
LITELLM_METADATA_ROUTES) request, so the key-presence check always picks
`litellm_metadata` there and silently reads no tags/identity at all.
Requiring the value to actually be a populated dict, matching
`_get_request_tags`'s own truthiness check in litellm_logging.py, only
ever prefers `litellm_metadata` when it is genuinely the field the proxy
wrote identity/tags into (LITELLM_METADATA_ROUTES pre-seed it before
admission runs, so it is always a populated dict by success time there)."""
litellm_metadata: Final = litellm_params_for_metadata.get("litellm_metadata")
if isinstance(litellm_metadata, Mapping) and litellm_metadata:
return "litellm_metadata"
return "metadata"
def extract_key_hash(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None:
"""Same single-authoritative-field lookup as
`model_based_tag_rate_limits_hook._extract_team_id`, but for the calling
virtual key's hash: `LiteLLMProxyRequestSetup` sets
`metadata["user_api_key"]` to `user_api_key_dict.api_key`, which despite
the plain name is already the hashed token (see `litellm_pre_call_utils.py`).
"""
active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING
key_hash: Final = active.get("user_api_key")
return key_hash if isinstance(key_hash, str) else None
def extract_key_alias(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None:
"""Same single-authoritative-field lookup as
`model_based_tag_rate_limits_hook._extract_team_id`, but for the calling
virtual key's own `key_alias`: `LiteLLMProxyRequestSetup` sets
`metadata["user_api_key_alias"]` to `user_api_key_dict.key_alias`
(see `litellm_pre_call_utils.py`)."""
active: Final = request_kwargs.get(metadata_variable_name) or EMPTY_MAPPING
key_alias: Final = active.get("user_api_key_alias")
return key_alias if isinstance(key_alias, str) else None
def fixed_length_identity(tag_value: str) -> str:
"""
`tag_value` is caller-controlled (whatever follows the tag_id prefix in
a caller-supplied tag) with no length or content bound. Embedding it
directly would let a caller inflate a hook's own in-memory dict keys
past what `max_in_memory_cache_size` bounds (that caps item *count*, not
key bytes) and grow unbounded Redis keys with no cap at all. Hashing to
a fixed-length digest bounds a hook's own contribution to key size
regardless of the caller's input, while still preserving distinctness
(two different tag values still resolve to two different buckets).
"""
return hashlib.sha256(tag_value.encode()).hexdigest()
def policy_fingerprint(entry: TagRateLimitEntry) -> str:
"""
Two entries can share a `name` and `tag_id` while genuinely disagreeing
on `limit`, `period_seconds`, `scope_by_key_hash`, or any of the scoping
fields -- both hooks already treat that as two distinct policies
elsewhere (model_based_tag_rate_limits_hook.py's own per-deployment
dedup signature is one example), so the Redis/in-memory bucket key must
too, or two differently-configured entries that happen to share a name
check and charge the identical counter. `scope_by_key_hash` specifically
needs its own slot here rather than relying on a hook's own `_hash_tag`
`key_hash`-derived suffix to carry it: that suffix is empty whenever
`key_hash` resolves to `None` (no virtual key on the call), which would
otherwise collide an unscoped entry with a key-hash-scoped one that
agrees on every other field. Hashed to a fixed-length digest for the
same reason `fixed_length_identity` hashes `tag_value`: an operator's
own `enabled_for`/`disabled_for`/`apply_to_key_alias`/`apply_to_models`
list has no length bound.
"""
fingerprint_source: Final = (
entry.limit,
entry.period_seconds,
entry.scope_by_key_hash,
scope_signature(entry.enabled_for),
scope_signature(entry.disabled_for),
entry.apply_to_key_alias,
entry.apply_to_models,
)
return hashlib.sha256(repr(fingerprint_source).encode()).hexdigest()[:16]
def bucket_ttl_seconds(entry: TagRateLimitEntry) -> int:
"""Redis (and in-memory fallback) TTL for a non-concurrency bucket key.
`entry.key_ttl_seconds` overrides the default of period_seconds + 3600
when set -- see TagRateLimitEntry.key_ttl_seconds."""
return entry.key_ttl_seconds if entry.key_ttl_seconds is not None else entry.period_seconds + 3600
# None => this entry shares the hook's single default cache partition
# (matching every entry's behavior before this override existed). Otherwise
# a value-stable signature -- not the override int alone -- so two different
# entries that happen to choose the identical max_in_memory_cache_size don't
# get merged into one shared partition; the same entry (same config content)
# always resolves to the same signature across config/index rebuilds, which
# is what keeps each hook's own `_partitions` cache from leaking a fresh
# partition every time its configuration is re-resolved.
# The str before the trailing int is `policy_fingerprint(entry)`: two
# entries can share tag_id/name while genuinely disagreeing on
# limit/period_seconds/scope_by_key_hash/enabled_for/disabled_for/
# apply_to_key_alias/apply_to_models -- both hooks already treat that as two
# distinct policies for bucket-key purposes, so a shared
# max_in_memory_cache_size must not route them onto the same partition
# either, or one entry's high-cardinality traffic can evict the other's
# active counters from a cache neither entry asked to share.
# max_in_memory_cache_size stays the trailing element: `partition_key[-1]`
# reads it directly to size the partition's cache.
PartitionKey: TypeAlias = tuple[str, str, str, int] | None
# Grouping type for async_log_success_event's per-partition tokens/dollars
# pipeline dispatch -- named only so the declaration fits on one line; see
# that method for why the grouping is needed.
PartitionOperations: TypeAlias = dict[PartitionKey, list[RedisPipelineIncrementOperation]]
def partition_key(entry: TagRateLimitEntry) -> PartitionKey:
if entry.max_in_memory_cache_size is None:
return None
return (
entry.tag_id,
entry.name,
policy_fingerprint(entry),
entry.max_in_memory_cache_size,
)

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

@ -0,0 +1,346 @@
"""
Unit tests for the primitives shared by both tag-scoped rate-limit hooks
(`model_based_tag_rate_limits_hook.py` and `global_tag_rate_limits_hook.py`).
These test the hook-independent logic in isolation: identity extraction,
`entry_applies` scoping, and the partition/bucket-TTL key helpers. Each
hook's own test file covers everything specific to how it wires these
primitives into its own admission/accounting engine.
"""
from litellm.proxy.hooks.tag_rate_limits_shared import (
bucket_ttl_seconds,
entry_applies,
extract_identity,
extract_key_hash,
fixed_length_identity,
partition_key,
)
from litellm.types.router import TagRateLimitEntry, TagRateLimitScope
# ---------------------------------------------------------------------------
# extract_identity
# ---------------------------------------------------------------------------
def test_extract_identity_matches_prefixed_tag():
assert extract_identity(["team_id:t1", "end_user_id:u1"], "end_user_id") == "u1"
def test_extract_identity_returns_none_when_absent():
assert extract_identity(["team_id:t1"], "end_user_id") is None
def test_extract_identity_skips_negation_tags():
"""A `!end_user_id:u1` routing-negation marker must never be read as identity."""
assert extract_identity(["!end_user_id:u1"], "end_user_id") is None
# ---------------------------------------------------------------------------
# fixed_length_identity -- tag_value is caller-controlled with no length
# bound; a hook's own contribution to a cache key must not grow with it
# ---------------------------------------------------------------------------
def test_fixed_length_identity_bounds_key_contribution_regardless_of_input_size():
"""
A caller can submit an arbitrarily long tag value (no length or content
bound is enforced upstream of either hook). Without hashing, that value
would go straight into an in-memory dict key (bypassing
max_in_memory_cache_size, which caps item *count* not key bytes) and an
unbounded-length Redis key (Redis has no key-count or key-size cap at
all here). A fixed-length digest bounds a hook's own contribution to
the key regardless of input size.
"""
huge_value = "x" * 5_000_000
digest = fixed_length_identity(huge_value)
assert len(digest) == 64 # sha256 hex digest length, independent of input size
def test_fixed_length_identity_preserves_distinctness():
"""Hashing must not collapse two different tag values onto one bucket."""
assert fixed_length_identity("user-a") != fixed_length_identity("user-b")
assert fixed_length_identity("user-a") == fixed_length_identity("user-a")
# ---------------------------------------------------------------------------
# extract_key_hash -- must read only the one field the server actually
# authenticates into, never fall back to the other
# ---------------------------------------------------------------------------
def test_extract_key_hash_ignores_a_forged_value_in_the_non_authoritative_field():
"""
On a route where litellm_metadata is authoritative, the server writes
the real hash there and never touches metadata -- so a caller-supplied
metadata.user_api_key must not be read at all, let alone win.
"""
request_kwargs = {
"metadata": {"user_api_key": "forged-by-caller"},
"litellm_metadata": {"user_api_key": "real-authenticated-hash"},
}
assert extract_key_hash(request_kwargs, "litellm_metadata") == "real-authenticated-hash"
def test_extract_key_hash_reads_metadata_when_it_is_the_authoritative_field():
request_kwargs = {"metadata": {"user_api_key": "real-hash"}}
assert extract_key_hash(request_kwargs, "metadata") == "real-hash"
# ---------------------------------------------------------------------------
# entry_applies -- enabled_for / disabled_for / apply_to_key_alias / apply_to_models
# ---------------------------------------------------------------------------
def test_entry_applies_with_none_of_the_scoping_fields_set():
entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is True
def test_entry_applies_disabled_for_on_its_own_tag_id_excludes_a_listed_value():
"""disabled_for's `tag_id` can be set to the entry's own tag_id, gating on
a subset of its own resolved identity rather than a second tag."""
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)),
)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is False
assert entry_applies(entry, ["end_user_id:u2"], None, None) is True
def test_entry_applies_enabled_for_on_its_own_tag_id_restricts_to_a_listed_value():
"""enabled_for's `tag_id` can likewise be set to the entry's own tag_id,
admitting only a hand-picked subset of its own resolved identity."""
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
enabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u2", "u3")),
)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is False
assert entry_applies(entry, ["end_user_id:u2"], None, None) is True
def test_entry_applies_matches_an_enabled_for_gate():
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)),
)
assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is True
def test_entry_applies_skips_when_enabled_for_gate_tag_is_absent():
"""
enabled_for is an allowlist gate: absence of the gate tag must not
satisfy it, unlike disabled_for below.
"""
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)),
)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is False
def test_entry_applies_skips_when_disabled_for_gate_matches():
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)),
)
assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False
def test_entry_applies_when_disabled_for_gate_tag_is_absent():
"""disabled_for is a denylist gate: absence of the gate tag has nothing
to match against, so the entry still applies."""
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)),
)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is True
def test_entry_applies_disabled_for_overrides_a_matching_enabled_for_gate():
"""Deny (disabled_for) takes effect independently of whether the
enabled_for gate itself matched, even when both target the same tag."""
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)),
disabled_for=TagRateLimitScope(tag_id="end_user_id", values=("u1",)),
)
assert entry_applies(entry, ["end_user_id:u1", "company_id:1032"], None, None) is False
def test_entry_applies_with_apply_to_key_alias_unset_applies_to_every_key():
entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400)
assert entry_applies(entry, ["end_user_id:u1"], "any-key-alias", None) is True
assert entry_applies(entry, ["end_user_id:u1"], None, None) is True
def test_entry_applies_admits_a_key_alias_on_the_allowlist():
entry = TagRateLimitEntry(
name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",)
)
assert entry_applies(entry, ["end_user_id:u1"], "team-a-key", None) is True
def test_entry_applies_rejects_a_key_alias_missing_from_the_allowlist():
entry = TagRateLimitEntry(
name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",)
)
assert entry_applies(entry, ["end_user_id:u1"], "team-b-key", None) is False
def test_entry_applies_rejects_when_key_has_no_alias_but_allowlist_is_set():
"""apply_to_key_alias is an allowlist gate: a key with no alias at all
never satisfies it, same as enabled_for's absent-gate-tag semantics."""
entry = TagRateLimitEntry(
name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_key_alias=("team-a-key",)
)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is False
def test_entry_applies_with_apply_to_models_unset_applies_to_every_model():
entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400)
assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True
assert entry_applies(entry, ["end_user_id:u1"], None, None) is True
def test_entry_applies_admits_a_model_on_the_apply_to_models_allowlist():
entry = TagRateLimitEntry(
name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",)
)
assert entry_applies(entry, ["end_user_id:u1"], None, "opus-chain") is True
def test_entry_applies_rejects_a_model_missing_from_the_apply_to_models_allowlist():
entry = TagRateLimitEntry(
name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",)
)
assert entry_applies(entry, ["end_user_id:u1"], None, "sonnet-chain") is False
def test_entry_applies_rejects_when_model_is_absent_but_apply_to_models_is_set():
"""apply_to_models is an allowlist gate: a request with no model at all
never satisfies it, same as apply_to_key_alias's absent-key semantics."""
entry = TagRateLimitEntry(
name="daily", tag_id="end_user_id", limit=500, period_seconds=86400, apply_to_models=("opus-chain",)
)
assert entry_applies(entry, ["end_user_id:u1"], None, None) is False
def test_entry_applies_apply_to_models_composes_with_apply_to_key_alias():
"""Both gates must pass: a request against the listed model but a
non-listed key alias must not apply, even though apply_to_models alone
would have admitted it."""
entry = TagRateLimitEntry(
name="daily",
tag_id="end_user_id",
limit=500,
period_seconds=86400,
apply_to_models=("opus-chain",),
apply_to_key_alias=("premium-key",),
)
assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "opus-chain") is True
assert entry_applies(entry, ["end_user_id:u1"], "other-key", "opus-chain") is False
assert entry_applies(entry, ["end_user_id:u1"], "premium-key", "sonnet-chain") is False
# ---------------------------------------------------------------------------
# partition_key -- entries that share max_in_memory_cache_size but disagree
# on any policy-fingerprinted field must never share a cache partition
# ---------------------------------------------------------------------------
def test_partition_key_distinguishes_entries_that_differ_only_by_scope_by_key_hash():
"""
scope_by_key_hash is part of the partition-key signature: two entries
identical in every other field but differing only on this flag are
different rate limits (different bucket keys per each hook's own
`_hash_tag`) and must never be routed to the same cache partition.
"""
unscoped = TagRateLimitEntry(
name="per_minute", tag_id="end_user_id", limit=5, period_seconds=60, max_in_memory_cache_size=100
)
scoped = TagRateLimitEntry(
name="per_minute",
tag_id="end_user_id",
limit=5,
period_seconds=60,
scope_by_key_hash=True,
max_in_memory_cache_size=100,
)
assert partition_key(unscoped) != partition_key(scoped)
def test_partition_key_distinguishes_entries_that_differ_only_by_scoping_fields():
"""
A plain, unscoped entry and a scoped override can legitimately share
name/tag_id/limit/period_seconds/scope_by_key_hash while disagreeing on
enabled_for/disabled_for/apply_to_key_alias/apply_to_models --
policy_fingerprint already treats that as two distinct policies, so a
shared max_in_memory_cache_size must not route them onto the same
in-memory partition either, or one entry's high-cardinality traffic can
evict the other's active counters from a cache neither entry asked to
share.
"""
base_kwargs = {
"name": "daily",
"tag_id": "end_user_id",
"limit": 100,
"period_seconds": 86400,
"max_in_memory_cache_size": 50,
}
unscoped = TagRateLimitEntry(**base_kwargs)
enabled_for_scoped = TagRateLimitEntry(
**base_kwargs, enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",))
)
disabled_for_scoped = TagRateLimitEntry(
**base_kwargs, disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",))
)
alias_scoped = TagRateLimitEntry(**base_kwargs, apply_to_key_alias=("premium-key",))
models_scoped = TagRateLimitEntry(**base_kwargs, apply_to_models=("opus-chain",))
keys = {
partition_key(unscoped),
partition_key(enabled_for_scoped),
partition_key(disabled_for_scoped),
partition_key(alias_scoped),
partition_key(models_scoped),
}
assert len(keys) == 5
# ---------------------------------------------------------------------------
# bucket_ttl_seconds -- per-tag Redis/bucket key TTL override
# ---------------------------------------------------------------------------
def test_bucket_ttl_seconds_defaults_to_period_plus_one_hour_when_unset():
entry = TagRateLimitEntry(name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60)
assert bucket_ttl_seconds(entry) == 60 + 3600
def test_bucket_ttl_seconds_honors_key_ttl_seconds_override():
entry = TagRateLimitEntry(
name="per_minute", tag_id="end_user_id", limit=1, period_seconds=60, key_ttl_seconds=120
)
assert bucket_ttl_seconds(entry) == 120