diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 56c9147e066..28da62e15f9 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -253,7 +253,10 @@ class InMemoryCache(BaseCache): ) -> list[float] | None: results: Final = [] for increment in increment_list: - result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) + # forward each op's own ttl; allow_ttl_override leaves an already-live ttl untouched + result = await self.async_increment( + increment["key"], increment["increment_value"], ttl=increment.get("ttl"), **kwargs + ) results.append(result) return results 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..75ae67dd120 --- /dev/null +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -0,0 +1,235 @@ +"""Primitives shared by both tag-scoped rate-limit hooks (model_based_tag_rate_limits_hook.py, +global_tag_rate_limits_hook.py): identity/scope extraction, policy fingerprinting, bucket-key +hashing, and cache-partitioning. Every name here is a genuine public export (no leading +underscore); each hook imports what it needs aliased back to its own historical private name.""" + +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") + +# requests/concurrency admit via atomic check-and-increment; tokens/dollars are only known +# after the response, so they stay a read-then-account-on-success check with an unavoidable 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, so call sites don't build a fresh {} +EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) + +# holds a strong reference to every fire-and-forget background task (concurrency release, +# token/dollar accounting) so the event loop's weak-ref-only tracking can't garbage-collect +# one mid-execution; each task's own completion callback discards its entry +BACKGROUND_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: see comment above + +# floor for a concurrency reservation's self-heal ttl, regardless of period_seconds, so an +# expiring-while-in-flight reservation can't silently admit past the limit +CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600 + +# single-key atomic check-and-increment (one key per call: every tag_rl key carries its own +# {..} hash tag, so a multi-key call could span shards and cross-slot error). refresh_ttl +# (ARGV[4]) distinguishes callers: "requests" is a fixed window whose ttl is set once and never +# extended; "concurrency" isn't windowed, so its crash-safety ttl must refresh on every admission. +TAG_RL_CHECK_AND_INCR_SCRIPT: Final = """ +local key = KEYS[1] +local limit = tonumber(ARGV[1]) +local increment = tonumber(ARGV[2]) +local ttl = tonumber(ARGV[3]) +local refresh_ttl = tonumber(ARGV[4]) +local current = tonumber(redis.call('GET', key) or 0) +if current + increment > limit then + return { 0, current } +end +local new_value = redis.call('INCRBY', key, increment) +if ttl > 0 then + if refresh_ttl == 1 then + redis.call('EXPIRE', key, ttl) + else + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl) + end + end +end +return { 1, new_value } +""" + +# atomic decrement floored at 0 (refund or concurrency release); floors via DEL rather than +# `SET key 0`, since a release on an already-expired key would otherwise recreate it with no ttl +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, used to fold enabled_for/disabled_for +# into a hashable form 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 hashable tuple for policy fingerprinting/dedup.""" + 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. `!`-prefixed tag-routing + negation markers are skipped so they're never 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 in order (deny before allow): disabled_for excludes a + matching value; enabled_for requires a matching value (absence fails, unlike disabled_for); + apply_to_models and apply_to_key_alias are allowlists an absent model/alias never satisfies. + An entry with none of these set always applies.""" + 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_authoritative_metadata_variable_name( + metadata_source: Mapping[str, object], +) -> Literal["metadata", "litellm_metadata"]: + """Picks metadata vs litellm_metadata by the unforgeable `user_api_key_auth` marker + `add_litellm_data_to_request` stamps into whichever bucket is actually authoritative -- + key presence or truthiness alone can be forged by a caller onto the wrong bucket.""" + litellm_metadata: Final = metadata_source.get("litellm_metadata") + if isinstance(litellm_metadata, Mapping) and "user_api_key_auth" in litellm_metadata: + return "litellm_metadata" + return "metadata" + + +def _active_metadata_bucket(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> Mapping[str, object]: + """request_kwargs carries metadata at its own top level at admission time, but only nested + under litellm_params by async_log_success_event/async_log_failure_event time; checks both.""" + top_level: Final = request_kwargs.get(metadata_variable_name) + if isinstance(top_level, Mapping): + return top_level + litellm_params: Final = request_kwargs.get("litellm_params") + if isinstance(litellm_params, Mapping): + nested: Final = litellm_params.get(metadata_variable_name) + if isinstance(nested, Mapping): + return nested + return EMPTY_MAPPING + + +def extract_key_hash(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None: + """Reads the calling virtual key's hash from the authoritative metadata bucket (nested under + litellm_params by log time, same as _active_metadata_bucket's other callers); `metadata["user_api_key"]` + is already the hashed token despite the plain name.""" + active: Final = _active_metadata_bucket(request_kwargs, metadata_variable_name) + 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: + """Reads the calling virtual key's own key_alias from the authoritative metadata bucket.""" + active: Final = _active_metadata_bucket(request_kwargs, metadata_variable_name) + 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, ...]: + """Puts server-computed `metadata.inherited_tags` ahead of caller-supplied tags, so a caller + can't submit e.g. `company_id:attacker-chosen` and shadow the calling key's real, same-prefix + tag; extract_identity/entry_applies both resolve tag_id via first-match-by-prefix.""" + active: Final = _active_metadata_bucket(request_kwargs, metadata_variable_name) + inherited_tags: Final = active.get("inherited_tags") + if not isinstance(inherited_tags, (list, tuple)) or not inherited_tags: + return tuple(tags) + return tuple(dict.fromkeys((*inherited_tags, *tags))) + + +def fixed_length_identity(tag_value: str) -> str: + """Hashes a caller-controlled, unbounded-length tag value to a fixed-length digest, bounding + a hook's own contribution to an in-memory cache key or Redis key regardless of input size.""" + return hashlib.sha256(tag_value.encode()).hexdigest() + + +def policy_fingerprint(entry: TagRateLimitEntry) -> str: + """Fixed-length digest folding every policy-distinguishing field of an entry (limit, + period_seconds, scope_by_key_hash, enabled_for/disabled_for, apply_to_key_alias/models) into + one bucket-key signature, so two differently-configured entries sharing a name never share + a counter.""" + 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.""" + 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. Otherwise a value-stable +# signature (policy_fingerprint(entry), not the override int alone) so two entries sharing a +# max_in_memory_cache_size don't merge into one partition; max_in_memory_cache_size stays the +# trailing element since `partition_key[-1]` reads it directly to size the partition's cache. +PartitionKey: TypeAlias = tuple[str, str, str, int] | None +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 fc09c40fe08..845ae38e06b 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, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -160,6 +161,126 @@ 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`, used by `TagRateLimitEntry.enabled_for`/`disabled_for`.""" + + tag_id: str + values: tuple[str, ...] + + model_config = ConfigDict(frozen=True) + + @model_validator(mode="after") + def _validate_tag_id(self) -> "TagRateLimitScope": + if not self.tag_id: + raise ValueError("tag_id must be a non-empty string") + return self + + @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+deduped so dedup-signature comparisons aren't order-sensitive; __setattr__ works around frozen=True + 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 the default bucket ttl (period_seconds + 3600); see _PROXY_ModelBasedTagRateLimitsHook._ttl_for + key_ttl_seconds: int | None = None + # overrides the shared litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200) with a dedicated partition + max_in_memory_cache_size: int | None = None + # gate this entry on another tag; disabled_for wins over enabled_for; an absent tag never matches either allowlist + enabled_for: TagRateLimitScope | None = None + disabled_for: TagRateLimitScope | None = None + apply_to_key_alias: tuple[str, ...] | None = None + apply_to_models: tuple[str, ...] | None = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _validate_tag_id(self) -> "TagRateLimitEntry": + if not self.tag_id: + raise ValueError("tag_id must be a non-empty string") + return self + + @model_validator(mode="after") + def _validate_limit(self) -> "TagRateLimitEntry": + # NaN compares False against every ordering operator, silently defeating whichever check gates this limit + if math.isnan(self.limit): + raise ValueError("limit must not be NaN") + if math.isinf(self.limit): + raise ValueError("limit must be finite") + 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 resets the counter early, letting it 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": + # same reason as TagRateLimitScope._normalize_values + 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. @@ -214,6 +335,8 @@ class ModelInfo(MirroredPricingParams): # in the spend log row's metadata. Set it on every deployment of the group. internal_router_model: 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/local_testing/test_tag_rate_limits_shared_redis.py b/tests/local_testing/test_tag_rate_limits_shared_redis.py new file mode 100644 index 00000000000..8c29cd288d7 --- /dev/null +++ b/tests/local_testing/test_tag_rate_limits_shared_redis.py @@ -0,0 +1,128 @@ +""" +Real-Redis execution coverage for tag_rate_limits_shared's two Lua scripts. + +Lives outside tests/test_litellm/ (which can only contain mocked tests, see +tests/test_litellm/readme.md) because fakeredis has no EVALSHA support without +the optional lupa dependency; these run against a throwaway local redis-server +instead (same idiom as tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py's +test_redis_lua_path_full_lifecycle). +""" + +import shutil +import socket +import subprocess +import time +from collections.abc import Iterator + +import pytest + +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.hooks.tag_rate_limits_shared import ( + TAG_RL_CHECK_AND_INCR_SCRIPT, + TAG_RL_DECR_FLOOR_ZERO_SCRIPT, +) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +@pytest.fixture +def redis_port() -> Iterator[int]: + if shutil.which("redis-server") is None: + pytest.skip("requires a local redis-server binary to exercise the Lua script path") + port = _free_port() + proc = subprocess.Popen( + ["redis-server", "--port", str(port), "--save", "", "--appendonly", "no"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + if sock.connect_ex(("127.0.0.1", port)) == 0: + break + else: + proc.terminate() + pytest.skip("local redis-server did not become ready in time") + yield port + finally: + proc.terminate() + proc.wait(timeout=5) + + +def _register(port: int, script: str): + return RedisCache(host="127.0.0.1", port=port).async_register_script(script) + + +@pytest.mark.asyncio +async def test_check_and_incr_admits_under_limit_and_sets_ttl(redis_port): + run = _register(redis_port, TAG_RL_CHECK_AND_INCR_SCRIPT) + admitted, new_value = await run(keys=["bucket1"], args=[5, 1, 60, 0]) + assert (admitted, new_value) == (1, 1) + + +@pytest.mark.asyncio +async def test_check_and_incr_rejects_over_limit_without_incrementing(redis_port): + run = _register(redis_port, TAG_RL_CHECK_AND_INCR_SCRIPT) + await run(keys=["bucket2"], args=[1, 1, 60, 0]) + rejected, current = await run(keys=["bucket2"], args=[1, 1, 60, 0]) + assert (rejected, current) == (0, 1) + + +@pytest.mark.asyncio +async def test_check_and_incr_requests_ttl_is_set_once_and_not_refreshed(redis_port): + """refresh_ttl=0 (the `requests` fixed-window semantics): the epoch-bucketed + TTL must be set on the first write and left alone after, or the bucket + outlives the epoch it's meant to reset at.""" + run = _register(redis_port, TAG_RL_CHECK_AND_INCR_SCRIPT) + await run(keys=["bucket3"], args=[5, 1, 100, 0]) + raw = RedisCache(host="127.0.0.1", port=redis_port) + client = raw.init_async_client() + first_ttl = await client.ttl("bucket3") + await run(keys=["bucket3"], args=[5, 1, 5, 0]) + second_ttl = await client.ttl("bucket3") + assert first_ttl > 5 + assert second_ttl > 5 # unchanged by the second call's much shorter ttl arg + + +@pytest.mark.asyncio +async def test_check_and_incr_concurrency_ttl_refreshes_on_every_admission(redis_port): + """refresh_ttl=1 (the `concurrency` reservation semantics): every admission + must push the crash-safety TTL back out, or a long-lived burst of traffic + expires the whole counter mid-flight.""" + run = _register(redis_port, TAG_RL_CHECK_AND_INCR_SCRIPT) + await run(keys=["bucket4"], args=[5, 1, 5, 1]) + await run(keys=["bucket4"], args=[5, 1, 100, 1]) + raw = RedisCache(host="127.0.0.1", port=redis_port) + client = raw.init_async_client() + ttl = await client.ttl("bucket4") + assert ttl > 5 + + +@pytest.mark.asyncio +async def test_decr_floor_zero_floors_at_zero_and_deletes_the_key(redis_port): + run_incr = _register(redis_port, TAG_RL_CHECK_AND_INCR_SCRIPT) + run_decr = _register(redis_port, TAG_RL_DECR_FLOOR_ZERO_SCRIPT) + await run_incr(keys=["bucket5"], args=[5, 1, 60, 1]) + + floored = await run_decr(keys=["bucket5"], args=[-5]) + assert floored == 0 + + raw = RedisCache(host="127.0.0.1", port=redis_port) + client = raw.init_async_client() + assert await client.get("bucket5") is None # floored via DEL, not a TTL-less `SET 0` + + +@pytest.mark.asyncio +async def test_decr_floor_zero_decrements_normally_when_result_stays_non_negative(redis_port): + run_incr = _register(redis_port, TAG_RL_CHECK_AND_INCR_SCRIPT) + run_decr = _register(redis_port, TAG_RL_DECR_FLOOR_ZERO_SCRIPT) + await run_incr(keys=["bucket6"], args=[5, 3, 60, 1]) + + remaining = await run_decr(keys=["bucket6"], args=[-1]) + assert remaining == 2 diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 40ad4f0c6f0..e7ac893bbdc 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -45,6 +45,39 @@ async def test_async_increment_delegates_to_locked_sync_path(): assert cache.get_cache("counter") == 5 +async def test_async_increment_pipeline_applies_each_operations_own_ttl(): + """ + Bugbot finding: async_increment_pipeline dropped each operation's own + "ttl" field, so a counter created through it always got the cache's + 600-second default_ttl regardless of what the caller actually configured + (e.g. an hourly or daily rate-limit window) -- the counter (and whatever + it was tracking against a limit) silently reset mid-window. + """ + cache = InMemoryCache(default_ttl=600) + await cache.async_increment_pipeline([{"key": "long-window-counter", "increment_value": 1, "ttl": 7200}]) + + ttl_remaining = await cache.async_get_ttl("long-window-counter") + assert ttl_remaining is not None + assert ttl_remaining > time.time() + 600 + + +async def test_async_increment_pipeline_preserves_an_existing_live_ttl_on_later_increments(): + """ + A second increment on the same still-live counter must not reset its + remaining ttl back up to the full window -- only the first increment + (the one that actually creates the counter) should set it. + """ + cache = InMemoryCache(default_ttl=600) + await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}]) + ttl_after_first = cache.ttl_dict["counter"] + + await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}]) + ttl_after_second = cache.ttl_dict["counter"] + + assert ttl_after_second == ttl_after_first + assert cache.get_cache("counter") == 2 + + def test_in_memory_openai_obj_cache(): from openai import OpenAI 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..a727d38e2eb --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limits_shared.py @@ -0,0 +1,446 @@ +""" +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_alias, + extract_key_hash, + fixed_length_identity, + order_tags_for_identity_resolution, + partition_key, + resolve_authoritative_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" + + +def test_extract_key_hash_ignores_a_non_mapping_authoritative_field(): + """Bugbot finding: metadata can arrive as an unparsed JSON string (see + apply_client_tag_policy_pre_auth's own docstring on multipart/extra_body + routes); a truthy non-Mapping must not reach .get() and crash.""" + request_kwargs = {"metadata": '{"user_api_key": "forged"}'} + assert extract_key_hash(request_kwargs, "metadata") is None + + +def test_extract_key_alias_ignores_a_non_mapping_authoritative_field(): + request_kwargs = {"metadata": '{"user_api_key_alias": "forged"}'} + assert extract_key_alias(request_kwargs, "metadata") is None + + +def test_extract_key_hash_finds_metadata_nested_under_litellm_params_at_log_time(): + """Bugbot finding: by async_log_success_event/async_log_failure_event time, + kwargs only carries metadata nested under litellm_params (see + _active_metadata_bucket's own docstring), not at request_kwargs' own top + level -- extract_key_hash must find it there too, or a key-hash-scoped + bucket's accounting silently reads a different key than admission did.""" + request_kwargs = {"litellm_params": {"metadata": {"user_api_key": "real-hash"}}} + assert extract_key_hash(request_kwargs, "metadata") == "real-hash" + + +def test_extract_key_alias_finds_metadata_nested_under_litellm_params_at_log_time(): + request_kwargs = {"litellm_params": {"metadata": {"user_api_key_alias": "real-alias"}}} + assert extract_key_alias(request_kwargs, "metadata") == "real-alias" + + +# --------------------------------------------------------------------------- +# resolve_authoritative_metadata_variable_name -- Veria AI finding: a +# caller-supplied, non-empty litellm_metadata must not be selected over +# metadata, the field the proxy actually wrote authenticated tags into +# --------------------------------------------------------------------------- + + +def test_resolve_authoritative_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.""" + metadata_source = {"litellm_metadata": {"marker": True}} + assert resolve_authoritative_metadata_variable_name(metadata_source) == "metadata" + + +def test_resolve_authoritative_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.""" + metadata_source = {"litellm_metadata": {"user_api_key_auth": object(), "tags": ["team_id:t1"]}} + assert resolve_authoritative_metadata_variable_name(metadata_source) == "litellm_metadata" + + +def test_resolve_authoritative_metadata_variable_name_defaults_to_metadata_when_litellm_metadata_absent(): + assert resolve_authoritative_metadata_variable_name({}) == "metadata" + + +def test_resolve_authoritative_metadata_variable_name_defaults_to_metadata_when_litellm_metadata_none(): + assert resolve_authoritative_metadata_variable_name({"litellm_metadata": None}) == "metadata" + + +def test_resolve_authoritative_metadata_variable_name_defaults_to_metadata_when_litellm_metadata_empty(): + assert resolve_authoritative_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 7fcb76b638f..d67a84d9d3c 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -9,6 +9,10 @@ from litellm.types.router import ( GenericLiteLLMParams, LiteLLM_Params, ModelInfo, + TagRateLimitEntry, + TagRateLimitGroup, + TagRateLimits, + TagRateLimitScope, ) from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams @@ -91,7 +95,7 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): - with pytest.raises(ValueError, match='validation error for ModelInfo'): + with pytest.raises(ValueError, match="validation error for ModelInfo"): ModelInfo(id="x", input_cost_per_token="free") @@ -146,3 +150,156 @@ def test_aws_session_tags_round_trip_as_sts_shaped_pairs(): def test_aws_session_tags_reject_shapes_sts_would_refuse(aws_session_tags): with pytest.raises(ValidationError, match="aws_session_tags"): LiteLLM_Params(model="bedrock/anthropic.claude-opus-5", aws_session_tags=aws_session_tags) + + +# 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) + + +def test_limit_nan_rejected(): + """NaN compares False against every ordering operator, so it would otherwise slip + past a `limit <= 0` style guard and silently defeat whichever check gates it.""" + with pytest.raises(ValueError, match="limit must not be NaN"): + TagRateLimitEntry(name="daily", limit=float("nan"), period_seconds=60) + + +def test_period_seconds_non_positive_rejected(): + with pytest.raises(ValueError, match="period_seconds must be a positive integer"): + TagRateLimitEntry(name="daily", limit=10, period_seconds=0) + + +def test_key_ttl_seconds_non_positive_rejected(): + with pytest.raises(ValueError, match="key_ttl_seconds must be a positive integer when set"): + TagRateLimitEntry(name="daily", limit=10, period_seconds=60, key_ttl_seconds=0) + + +def test_max_in_memory_cache_size_non_positive_rejected(): + with pytest.raises(ValueError, match="max_in_memory_cache_size must be a positive integer"): + TagRateLimitEntry(name="daily", limit=10, period_seconds=60, max_in_memory_cache_size=0) + + +def test_apply_to_key_alias_empty_list_rejected(): + with pytest.raises(ValueError, match="apply_to_key_alias must be a non-empty list"): + TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_key_alias=()) + + +def test_apply_to_key_alias_sorted_and_deduped(): + entry = TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_key_alias=("b", "a", "a")) + assert entry.apply_to_key_alias == ("a", "b") + + +def test_apply_to_models_empty_list_rejected(): + with pytest.raises(ValueError, match="apply_to_models must be a non-empty list"): + TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_models=()) + + +def test_apply_to_models_sorted_and_deduped(): + entry = TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_models=("gpt-4o", "claude", "claude")) + assert entry.apply_to_models == ("claude", "gpt-4o") + + +def test_scope_tag_id_empty_string_rejected(): + """An empty tag_id makes identity lookup search for a bare `:` prefix, silently + never matching instead of erroring at config load time (Greptile P1 on #39902).""" + with pytest.raises(ValueError, match="tag_id must be a non-empty string"): + TagRateLimitScope(tag_id="", values=("1032",)) + + +def test_scope_values_empty_list_rejected(): + with pytest.raises(ValueError, match="values must be a non-empty list"): + TagRateLimitScope(tag_id="company_id", values=()) + + +def test_scope_values_sorted_and_deduped(): + scope = TagRateLimitScope(tag_id="company_id", values=("1032", "1001", "1001")) + assert scope.values == ("1001", "1032") + + +def test_scope_is_frozen(): + scope = TagRateLimitScope(tag_id="company_id", values=("1032",)) + with pytest.raises(ValidationError): + scope.tag_id = "other_tag" + + +def test_entry_tag_id_empty_string_rejected(): + """Same silent-no-match failure mode as TagRateLimitScope.tag_id (Greptile P1 on + #39902); the default is non-empty, but an explicit override could still be empty.""" + with pytest.raises(ValueError, match="tag_id must be a non-empty string"): + TagRateLimitEntry(name="daily", tag_id="", limit=10, period_seconds=60) + + +def test_entry_accepts_enabled_for_and_disabled_for_scopes(): + entry = TagRateLimitEntry( + name="daily", + limit=10, + period_seconds=60, + enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)), + disabled_for=TagRateLimitScope(tag_id="company_id", values=("9999",)), + ) + assert entry.enabled_for.values == ("1032",) + assert entry.disabled_for.values == ("9999",) + + +def test_group_defaults_to_no_limits(): + assert TagRateLimitGroup().limits == () + + +def test_rate_limits_group_holds_multiple_entries(): + entries = ( + TagRateLimitEntry(name="daily", limit=10, period_seconds=60), + TagRateLimitEntry(name="weekly", limit=100, period_seconds=604800), + ) + group = TagRateLimitGroup(limits=entries) + assert group.limits == entries + + +def test_model_info_tag_rate_limits_defaults_to_none(): + assert ModelInfo(id="x").tag_rate_limits is None + + +def test_model_info_accepts_tag_rate_limits(): + limits = TagRateLimits( + token_limits=TagRateLimitGroup(limits=(TagRateLimitEntry(name="daily", limit=10, period_seconds=60),)) + ) + info = ModelInfo(id="x", tag_rate_limits=limits) + assert info.tag_rate_limits.token_limits.limits[0].name == "daily" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7ed8df6815d..6d89d8d3ec5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36665,6 +36665,60 @@ 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`, used by `TagRateLimitEntry.enabled_for`/`disabled_for`. + */ + 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 @@ -39656,6 +39710,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 */