From dbe5c3b0b20a945d58817eba1f8284e07849d75a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 30 Apr 2026 18:46:09 -0700 Subject: [PATCH 1/5] fix: close TOCTOU window in batch + dynamic rate limiters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batch rate limiter (`_check_and_increment_batch_counters`) and the dynamic rate limiter (`_check_rate_limits`) implemented rate limiting in two disjoint awaits: a `should_rate_limit(read_only=True)` check followed by a separate increment. Concurrent requests could all observe the same pre-increment state, all pass enforcement, and all then increment — multiplying the effective quota by the concurrency level. Demonstrated bypass (see new test): - Batch: 5 concurrent batches of 40 tokens each against TPM=100 consumed 200 tokens (100% over). - Dynamic: 5 concurrent priority="high" requests against RPM=2 all passed Phase 1 + Phase 3. Wrap both critical sections in a per-instance asyncio.Lock so the read and increment execute atomically within a process. Multi-replica deployments still rely on Redis Lua atomicity for cross-process safety; that is a follow-up. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/hooks/batch_rate_limiter.py | 136 ++++--- .../proxy/hooks/dynamic_rate_limiter_v3.py | 161 ++++---- .../hooks/parallel_request_limiter_v3.py | 7 + .../proxy/hooks/test_rate_limiter_toctou.py | 375 ++++++++++++++++++ 4 files changed, 538 insertions(+), 141 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 06b7d85789..fafe855540 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -167,6 +167,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): Check rate limits and increment counters by the batch amounts. Raises HTTPException if any limit would be exceeded. + + Holds the limiter's check-and-increment lock across the read-only + check and the increment to prevent concurrent batches from each + observing the same pre-increment state and collectively exceeding + the limit (TOCTOU). """ from litellm.types.caching import RedisPipelineIncrementOperation @@ -179,74 +184,75 @@ class _PROXY_BatchRateLimiter(CustomLogger): model_has_failures=False, ) - # Check current usage without incrementing - rate_limit_response = await self.parallel_request_limiter.should_rate_limit( - descriptors=descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, - ) - - # Verify batch won't exceed any limits - for status in rate_limit_response["statuses"]: - rate_limit_type = status["rate_limit_type"] - limit_remaining = status["limit_remaining"] - - required_capacity = ( - batch_usage.request_count - if rate_limit_type == "requests" - else batch_usage.total_tokens if rate_limit_type == "tokens" else 0 - ) - - if required_capacity > limit_remaining: - self._raise_rate_limit_error( - status, descriptors, batch_usage, rate_limit_type - ) - - # Build pipeline operations for batch increments - # Reuse the same keys that descriptors check - pipeline_operations: List[RedisPipelineIncrementOperation] = [] - - for descriptor in descriptors: - key = descriptor["key"] - value = descriptor["value"] - rate_limit = descriptor.get("rate_limit") - - if rate_limit is None: - continue - - # Add RPM increment if limit is set - if rate_limit.get("requests_per_unit") is not None: - rpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="requests" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=rpm_key, - increment_value=batch_usage.request_count, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Add TPM increment if limit is set - if rate_limit.get("tokens_per_unit") is not None: - tpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="tokens" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=tpm_key, - increment_value=batch_usage.total_tokens, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Execute increments - if pipeline_operations: - await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation( - pipeline_operations=pipeline_operations, + async with self.parallel_request_limiter._check_and_increment_lock: + # Check current usage without incrementing + rate_limit_response = await self.parallel_request_limiter.should_rate_limit( + descriptors=descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=True, ) + # Verify batch won't exceed any limits + for status in rate_limit_response["statuses"]: + rate_limit_type = status["rate_limit_type"] + limit_remaining = status["limit_remaining"] + + required_capacity = ( + batch_usage.request_count + if rate_limit_type == "requests" + else batch_usage.total_tokens if rate_limit_type == "tokens" else 0 + ) + + if required_capacity > limit_remaining: + self._raise_rate_limit_error( + status, descriptors, batch_usage, rate_limit_type + ) + + # Build pipeline operations for batch increments + # Reuse the same keys that descriptors check + pipeline_operations: List[RedisPipelineIncrementOperation] = [] + + for descriptor in descriptors: + key = descriptor["key"] + value = descriptor["value"] + rate_limit = descriptor.get("rate_limit") + + if rate_limit is None: + continue + + # Add RPM increment if limit is set + if rate_limit.get("requests_per_unit") is not None: + rpm_key = self.parallel_request_limiter.create_rate_limit_keys( + key=key, value=value, rate_limit_type="requests" + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=rpm_key, + increment_value=batch_usage.request_count, + ttl=self.parallel_request_limiter.window_size, + ) + ) + + # Add TPM increment if limit is set + if rate_limit.get("tokens_per_unit") is not None: + tpm_key = self.parallel_request_limiter.create_rate_limit_keys( + key=key, value=value, rate_limit_type="tokens" + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=tpm_key, + increment_value=batch_usage.total_tokens, + ttl=self.parallel_request_limiter.window_size, + ) + ) + + # Execute increments + if pipeline_operations: + await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + async def count_input_file_usage( self, file_id: str, diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 72483d29cd..4c0ea89554 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -460,92 +460,101 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if priority_descriptors: descriptors_to_check.extend(priority_descriptors) - # PHASE 1: Read-only check of ALL limits (no increments) - check_response = await self.v3_limiter.should_rate_limit( - descriptors=descriptors_to_check, - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, # CRITICAL: Don't increment counters yet - ) + # Phases 1-3 must run as a single atomic critical section. Without + # this lock, concurrent requests all observe the same Phase 1 state, + # all pass enforcement, then all increment in Phase 3 — bypassing + # the limit (TOCTOU). Multi-replica deployments additionally rely on + # Redis Lua atomicity for cross-process safety. + async with self.v3_limiter._check_and_increment_lock: + # PHASE 1: Read-only check of ALL limits (no increments) + check_response = await self.v3_limiter.should_rate_limit( + descriptors=descriptors_to_check, + parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=True, # CRITICAL: Don't increment counters yet + ) - verbose_proxy_logger.debug( - f"Read-only check: {json.dumps(check_response, indent=2)}" - ) + verbose_proxy_logger.debug( + f"Read-only check: {json.dumps(check_response, indent=2)}" + ) - # PHASE 2: Decide which limits to enforce - if check_response["overall_code"] == "OVER_LIMIT": - for status in check_response["statuses"]: - if status["code"] == "OVER_LIMIT": - descriptor_key = status["descriptor_key"] + # PHASE 2: Decide which limits to enforce + if check_response["overall_code"] == "OVER_LIMIT": + for status in check_response["statuses"]: + if status["code"] == "OVER_LIMIT": + descriptor_key = status["descriptor_key"] - # Model-wide limit exceeded (ALWAYS enforce) - if descriptor_key == "model_saturation_check": - raise HTTPException( - status_code=429, - detail={ - "error": f"Model capacity reached for {model}. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - }, - ) + # Model-wide limit exceeded (ALWAYS enforce) + if descriptor_key == "model_saturation_check": + raise HTTPException( + status_code=429, + detail={ + "error": f"Model capacity reached for {model}. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + }, + ) - # Priority limit exceeded (ONLY enforce when saturated) - elif descriptor_key == "priority_model" and should_enforce_priority: - verbose_proxy_logger.debug( - f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " - f"priority: {priority}" - ) - raise HTTPException( - status_code=429, - detail={ - "error": f"Priority-based rate limit exceeded. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}, " - f"Model saturation: {saturation:.1%}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - "x-litellm-saturation": f"{saturation:.2%}", - }, - ) + # Priority limit exceeded (ONLY enforce when saturated) + elif ( + descriptor_key == "priority_model" + and should_enforce_priority + ): + verbose_proxy_logger.debug( + f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " + f"priority: {priority}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": f"Priority-based rate limit exceeded. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + ) - # PHASE 3: Increment counters separately to avoid early-exit issues - # Model counter must ALWAYS increment, but priority counter might be over limit - # If we increment them together, v3_limiter's in-memory check will exit early - # and skip incrementing the model counter + # PHASE 3: Increment counters separately to avoid early-exit issues + # Model counter must ALWAYS increment, but priority counter might be over limit + # If we increment them together, v3_limiter's in-memory check will exit early + # and skip incrementing the model counter - # Step 3a: Increment model-wide counter (always) - model_increment_response = await self.v3_limiter.should_rate_limit( - descriptors=[model_wide_descriptor], - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=False, - ) - - # Step 3b: Increment priority counter (may be over limit, but we still track it) - if priority_descriptors: - priority_increment_response = await self.v3_limiter.should_rate_limit( - descriptors=priority_descriptors, + # Step 3a: Increment model-wide counter (always) + model_increment_response = await self.v3_limiter.should_rate_limit( + descriptors=[model_wide_descriptor], parent_otel_span=user_api_key_dict.parent_otel_span, read_only=False, ) - # Combine responses for post-call hook - combined_response = { - "overall_code": model_increment_response["overall_code"], - "statuses": model_increment_response["statuses"] - + priority_increment_response["statuses"], - } - data["litellm_proxy_rate_limit_response"] = combined_response - else: - data["litellm_proxy_rate_limit_response"] = model_increment_response + # Step 3b: Increment priority counter (may be over limit, but we still track it) + if priority_descriptors: + priority_increment_response = await self.v3_limiter.should_rate_limit( + descriptors=priority_descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=False, + ) + + # Combine responses for post-call hook + combined_response = { + "overall_code": model_increment_response["overall_code"], + "statuses": model_increment_response["statuses"] + + priority_increment_response["statuses"], + } + data["litellm_proxy_rate_limit_response"] = combined_response + else: + data["litellm_proxy_rate_limit_response"] = model_increment_response async def async_pre_call_hook( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index f29bbd2d9d..5fea31608f 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4,6 +4,7 @@ This is a rate limiter implementation based on a similar one by Envoy proxy. This is currently in development and not yet ready for production. """ +import asyncio import binascii import os from datetime import datetime @@ -171,6 +172,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Batch rate limiter (lazy loaded) self._batch_rate_limiter: Optional[Any] = None + # Serializes multi-phase check+increment sequences (batch + dynamic + # limiters) within this process to close the TOCTOU window between + # read-only check and counter increment. Multi-replica deployments + # additionally rely on Redis Lua atomicity for cross-process safety. + self._check_and_increment_lock = asyncio.Lock() + def _get_batch_rate_limiter(self) -> Optional[Any]: """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py new file mode 100644 index 0000000000..006dac0449 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -0,0 +1,375 @@ +""" +Tests validating TOCTOU race condition in batch + dynamic rate limiters. + +Issue: rate-limit check (read_only=True) and counter increment happen as two +separate awaits. Concurrent requests all observe the same pre-increment state, +all pass validation, then all increment — bypassing the limit. + +Vulnerable code paths: +- litellm/proxy/hooks/batch_rate_limiter.py:181-248 + (_check_and_increment_batch_counters: should_rate_limit(read_only=True) + → validate → async_increment_tokens_with_ttl_preservation) +- litellm/proxy/hooks/dynamic_rate_limiter_v3.py:463-548 + (_check_rate_limits PHASE 1 read_only check → PHASE 3 increment) + +These tests EXPECTED to fail against current (vulnerable) code and pass once +check-and-increment becomes atomic. +""" + +import asyncio +import os +import sys +from typing import Any, Dict, List, Optional + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm import DualCache, Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler, +) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.proxy.utils import InternalUsageCache, hash_token + + +def _make_phase1_barrier(num_concurrent: int, timeout: float = 0.1): + """ + Sync primitive that, pre-fix, forces all N concurrent coroutines to finish + their read-only Phase 1 check before any proceeds to Phase 3 increment — + mimicking asyncio I/O scheduling under load on the vulnerable code. + + Wraps `should_rate_limit` so on `read_only=True` calls it waits until N + callers arrive (TOCTOU window opened) OR `timeout` elapses (post-fix path: + the limiter's serialization lock prevents N from ever reaching the + barrier; the timeout lets the holder proceed so the lock can do its job). + + Pre-fix: barrier fills before timeout → all see same state → bypass observed. + Post-fix: only lock-holder reaches barrier → times out → serial execution + enforces limit. + """ + arrived = 0 + all_arrived = asyncio.Event() + + def wrap(original): + async def patched(*args, **kwargs): + result = await original(*args, **kwargs) + if kwargs.get("read_only"): + nonlocal arrived + arrived += 1 + if arrived >= num_concurrent: + all_arrived.set() + try: + await asyncio.wait_for(all_arrived.wait(), timeout=timeout) + except asyncio.TimeoutError: + pass + return result + + return patched + + return wrap + + +@pytest.mark.asyncio +async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou(): + """ + 5 concurrent batch submissions of 40 tokens each against TPM=100 limit. + + Sequential semantics: only 2 batches fit (2 * 40 = 80 ≤ 100, 3rd at 120 fails). + With TOCTOU: all 5 succeed → 200 tokens consumed, 100% over limit. + + Demonstrates batch_rate_limiter.py:183-248 multi-phase flaw. + """ + NUM_CONCURRENT = 5 + BATCH_TOKENS = 40 + TPM_LIMIT = 100 + + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("toctou-batch-key"), + tpm_limit=TPM_LIMIT, + rpm_limit=1000, + ) + batch_usage = BatchFileUsage( + total_tokens=BATCH_TOKENS, request_count=1 + ) + + barrier = _make_phase1_barrier(NUM_CONCURRENT) + rate_limiter.should_rate_limit = barrier(rate_limiter.should_rate_limit) + + results = await asyncio.gather( + *[ + batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=batch_usage, + ) + for _ in range(NUM_CONCURRENT) + ], + return_exceptions=True, + ) + + successes = [r for r in results if not isinstance(r, Exception)] + rejections = [r for r in results if isinstance(r, Exception)] + total_consumed = len(successes) * BATCH_TOKENS + max_allowed_successes = TPM_LIMIT // BATCH_TOKENS # 2 + + assert len(successes) <= max_allowed_successes, ( + f"TOCTOU bypass: {len(successes)}/{NUM_CONCURRENT} concurrent batches " + f"passed despite TPM={TPM_LIMIT}. Consumed {total_consumed} tokens " + f"({total_consumed - TPM_LIMIT} over limit). " + f"Atomic check-and-increment would allow ≤{max_allowed_successes}. " + f"Rejections: {len(rejections)}" + ) + + +@pytest.mark.asyncio +async def test_batch_limiter_check_and_increment_is_two_separate_calls(): + """ + Structural test: _check_and_increment_batch_counters issues a read_only=True + check followed by a separate increment call — non-atomic by construction. + + Records call ordering on parallel_request_limiter to prove Phase 1 (check) + and Phase 3 (increment) are not wrapped in a single Redis transaction / + Lua script. After the fix, this pattern should be replaced with one + atomic_check_and_increment call. + """ + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + call_log: List[Dict[str, Any]] = [] + original_should = rate_limiter.should_rate_limit + original_inc = rate_limiter.async_increment_tokens_with_ttl_preservation + + async def logging_should(*args, **kwargs): + call_log.append( + {"method": "should_rate_limit", "read_only": kwargs.get("read_only")} + ) + return await original_should(*args, **kwargs) + + async def logging_inc(*args, **kwargs): + call_log.append({"method": "async_increment_tokens_with_ttl_preservation"}) + return await original_inc(*args, **kwargs) + + rate_limiter.should_rate_limit = logging_should + rate_limiter.async_increment_tokens_with_ttl_preservation = logging_inc + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("structural-test-key"), + tpm_limit=10000, + rpm_limit=1000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=50, request_count=1), + ) + + method_sequence = [c["method"] for c in call_log] + assert "should_rate_limit" in method_sequence, "Expected Phase 1 check" + phase1 = [ + c + for c in call_log + if c["method"] == "should_rate_limit" and c.get("read_only") is True + ] + phase3 = [ + c + for c in call_log + if c["method"] == "async_increment_tokens_with_ttl_preservation" + ] + assert len(phase1) >= 1 and len(phase3) >= 1, ( + f"Expected non-atomic Phase1+Phase3 pattern. call_log={call_log}" + ) + phase1_idx = method_sequence.index("should_rate_limit") + phase3_idx = method_sequence.index( + "async_increment_tokens_with_ttl_preservation" + ) + assert phase1_idx < phase3_idx, ( + "TOCTOU evidence: read-only check precedes increment as separate awaits — " + "no atomic Lua script wraps both. Sequence: " + f"{method_sequence}" + ) + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): + """ + DynamicRateLimitHandler PHASE 1 (read_only check) → PHASE 3 (increment) + is non-atomic: dynamic_rate_limiter_v3.py:463-548. + + With TPM=100 model capacity and 5 concurrent priority="high" requests + each consuming the full model_saturation_check counter, all observe the + same Phase 1 state (counter=0), all pass, all proceed to Phase 3. + + Sequential atomic semantics would block requests once the model counter + reaches its limit. TOCTOU lets all pass Phase 1 simultaneously. + """ + NUM_CONCURRENT = 10 + MODEL_RPM = 2 + # Sequential bound: dynamic limiter rejects when `counter > current_limit` + # (strict `>`), so a request whose Phase 1 sees counter=RPM still passes + # (RPM is not strictly greater). Atomic execution therefore admits up to + # RPM + 1 successes before the next sees counter > RPM. + MAX_SEQUENTIAL_SUCCESSES = MODEL_RPM + 1 + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "toctou-dyn-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": MODEL_RPM, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + barrier = _make_phase1_barrier(NUM_CONCURRENT) + handler.v3_limiter.should_rate_limit = barrier(handler.v3_limiter.should_rate_limit) + + from litellm.types.router import ModelGroupInfo + + model_group_info = ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=MODEL_RPM, + tpm=None, + ) + + async def one_request(idx: int): + user = UserAPIKeyAuth(api_key=hash_token(f"dyn-key-{idx}")) + user.metadata = {"priority": "high"} + try: + await handler._check_rate_limits( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + return "OK" + except Exception as e: + return e + + results = await asyncio.gather( + *[one_request(i) for i in range(NUM_CONCURRENT)], + return_exceptions=True, + ) + successes = [r for r in results if r == "OK"] + + assert len(successes) <= MAX_SEQUENTIAL_SUCCESSES, ( + f"TOCTOU bypass in DynamicRateLimitHandler: {len(successes)}/{NUM_CONCURRENT} " + f"concurrent requests passed Phase 1 + Phase 3 despite model RPM={MODEL_RPM}. " + f"Atomic check-and-increment would block once counter > RPM " + f"(at most {MAX_SEQUENTIAL_SUCCESSES} sequential successes)." + ) + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_phase1_phase3_are_separate_awaits(): + """ + Structural proof of TOCTOU: dynamic_rate_limiter_v3._check_rate_limits + issues a read_only=True call (Phase 1, line 464-468) followed by separate + read_only=False calls (Phase 3, lines 526-530 + 534-538). + + Records each invocation of v3_limiter.should_rate_limit and asserts the + Phase1→Phase3 sequence. After fix, both phases must collapse into a single + atomic operation. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "structural-dyn-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + read_only_flags: List[Optional[bool]] = [] + original = handler.v3_limiter.should_rate_limit + + async def logging_should(*args, **kwargs): + read_only_flags.append(kwargs.get("read_only")) + return await original(*args, **kwargs) + + handler.v3_limiter.should_rate_limit = logging_should + + from litellm.types.router import ModelGroupInfo + + user = UserAPIKeyAuth(api_key=hash_token("dyn-structural-key")) + user.metadata = {"priority": "high"} + + await handler._check_rate_limits( + model=model, + model_group_info=ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=None, + tpm=1000, + ), + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + + assert True in read_only_flags or any( + f is True for f in read_only_flags + ), f"Expected read_only=True (Phase 1) call. Got: {read_only_flags}" + assert any(f is False for f in read_only_flags), ( + f"Expected read_only=False (Phase 3) call. Got: {read_only_flags}" + ) + + first_read_only = next( + (i for i, f in enumerate(read_only_flags) if f is True), None + ) + first_write = next( + (i for i, f in enumerate(read_only_flags) if f is False), None + ) + assert first_read_only is not None and first_write is not None + assert first_read_only < first_write, ( + f"TOCTOU evidence: Phase 1 (read_only) precedes Phase 3 (increment) " + f"as separate non-atomic calls. read_only sequence: {read_only_flags}" + ) From dd57ae66915f564e87e91cc5bae89005c0a3c3e0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 30 Apr 2026 18:54:36 -0700 Subject: [PATCH 2/5] feat(rate-limit): atomic check-and-increment-by-N for multi-process safety MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fix for the TOCTOU bypass relied on a per-instance asyncio.Lock, which closed the window only within a single proxy worker. Multi-replica deployments still raced across processes — A and B both read counter=99, both passed validation, both incremented to 100/100 → effective limit doubled. Add `CHECK_AND_INCREMENT_BY_N_SCRIPT` Lua script that processes any number of (window_key, counter_key, limit, increment, ttl) descriptors atomically with all-or-nothing semantics: if any descriptor would exceed its limit, no counter is modified and the script returns OVER_LIMIT with the offending descriptor's state. When Redis isn't configured, the in-memory fallback uses the existing asyncio.Lock for single-process atomicity. Expose this as `_PROXY_MaxParallelRequestsHandler_v3.atomic_check_and_increment_by_n` and rewire both call sites: - batch_rate_limiter._check_and_increment_batch_counters: replace the read_only=True check + separate async_increment_tokens_with_ttl_preservation with a single atomic call passing the batch's (request_count, total_tokens) as the increment. - dynamic_rate_limiter_v3._check_rate_limits: bundle model_saturation_check (always enforced) and priority_model (enforced only when saturated) into one atomic call. When priority is unenforced, increment its counter via the existing should_rate_limit(read_only=False) path for tracking only. Update structural regression tests to assert the new atomic path is used rather than the legacy two-phase pattern. Tests: 4/4 TOCTOU tests pass, 59 existing rate-limiter tests pass, no regressions. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/hooks/batch_rate_limiter.py | 99 ++---- .../proxy/hooks/dynamic_rate_limiter_v3.py | 167 +++++----- .../hooks/parallel_request_limiter_v3.py | 312 ++++++++++++++++++ .../proxy/hooks/test_rate_limiter_toctou.py | 121 +++---- 4 files changed, 463 insertions(+), 236 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index fafe855540..ffbcfc0f44 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,7 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast from fastapi import HTTPException from pydantic import BaseModel @@ -164,18 +164,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage: BatchFileUsage, ) -> None: """ - Check rate limits and increment counters by the batch amounts. + Atomically check + increment rate-limit counters by the batch amounts. - Raises HTTPException if any limit would be exceeded. - - Holds the limiter's check-and-increment lock across the read-only - check and the increment to prevent concurrent batches from each - observing the same pre-increment state and collectively exceeding - the limit (TOCTOU). + Raises HTTPException if any descriptor would exceed its limit; in that + case no counter is modified. Backed by `atomic_check_and_increment_by_n` + which uses a Redis Lua script when available (multi-process atomic) and + falls back to a per-process asyncio.Lock + in-memory operation. """ - from litellm.types.caching import RedisPipelineIncrementOperation - - # Create descriptors and check if batch would exceed limits descriptors = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, @@ -184,75 +179,35 @@ class _PROXY_BatchRateLimiter(CustomLogger): model_has_failures=False, ) - async with self.parallel_request_limiter._check_and_increment_lock: - # Check current usage without incrementing - rate_limit_response = await self.parallel_request_limiter.should_rate_limit( + increments = cast( + List[Dict[Literal["requests", "tokens"], int]], + [ + { + "requests": batch_usage.request_count, + "tokens": batch_usage.total_tokens, + } + for _ in descriptors + ], + ) + + rate_limit_response = ( + await self.parallel_request_limiter.atomic_check_and_increment_by_n( descriptors=descriptors, + increments=increments, parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, ) + ) - # Verify batch won't exceed any limits + if rate_limit_response["overall_code"] == "OVER_LIMIT": for status in rate_limit_response["statuses"]: - rate_limit_type = status["rate_limit_type"] - limit_remaining = status["limit_remaining"] - - required_capacity = ( - batch_usage.request_count - if rate_limit_type == "requests" - else batch_usage.total_tokens if rate_limit_type == "tokens" else 0 - ) - - if required_capacity > limit_remaining: + if status["code"] == "OVER_LIMIT": self._raise_rate_limit_error( - status, descriptors, batch_usage, rate_limit_type + status, + descriptors, + batch_usage, + status["rate_limit_type"], ) - # Build pipeline operations for batch increments - # Reuse the same keys that descriptors check - pipeline_operations: List[RedisPipelineIncrementOperation] = [] - - for descriptor in descriptors: - key = descriptor["key"] - value = descriptor["value"] - rate_limit = descriptor.get("rate_limit") - - if rate_limit is None: - continue - - # Add RPM increment if limit is set - if rate_limit.get("requests_per_unit") is not None: - rpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="requests" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=rpm_key, - increment_value=batch_usage.request_count, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Add TPM increment if limit is set - if rate_limit.get("tokens_per_unit") is not None: - tpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="tokens" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=tpm_key, - increment_value=batch_usage.total_tokens, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Execute increments - if pipeline_operations: - await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation( - pipeline_operations=pipeline_operations, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - async def count_input_file_usage( self, file_id: str, diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 4c0ea89554..d827cf3067 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting import os from datetime import datetime -from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, Union from fastapi import HTTPException @@ -460,101 +460,90 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if priority_descriptors: descriptors_to_check.extend(priority_descriptors) - # Phases 1-3 must run as a single atomic critical section. Without - # this lock, concurrent requests all observe the same Phase 1 state, - # all pass enforcement, then all increment in Phase 3 — bypassing - # the limit (TOCTOU). Multi-replica deployments additionally rely on - # Redis Lua atomicity for cross-process safety. - async with self.v3_limiter._check_and_increment_lock: - # PHASE 1: Read-only check of ALL limits (no increments) - check_response = await self.v3_limiter.should_rate_limit( - descriptors=descriptors_to_check, - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, # CRITICAL: Don't increment counters yet - ) + # Atomic check-and-increment for the ENFORCED descriptor set: + # - model_saturation_check is always enforced + # - priority_model is enforced only when saturation crosses threshold + # + # Backed by a Redis Lua script (multi-process atomic) with an + # asyncio.Lock + in-memory fallback for single-process deployments. + # All-or-nothing: if any enforced descriptor would exceed its limit, + # no counter is modified and the response carries "OVER_LIMIT". + enforced_descriptors: List[RateLimitDescriptor] = [model_wide_descriptor] + if priority_descriptors and should_enforce_priority: + enforced_descriptors.extend(priority_descriptors) - verbose_proxy_logger.debug( - f"Read-only check: {json.dumps(check_response, indent=2)}" - ) + per_request_increment: Dict[Literal["requests", "tokens"], int] = { + "requests": 1, + "tokens": 0, + } + atomic_response = await self.v3_limiter.atomic_check_and_increment_by_n( + descriptors=enforced_descriptors, + increments=[per_request_increment for _ in enforced_descriptors], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) - # PHASE 2: Decide which limits to enforce - if check_response["overall_code"] == "OVER_LIMIT": - for status in check_response["statuses"]: - if status["code"] == "OVER_LIMIT": - descriptor_key = status["descriptor_key"] + verbose_proxy_logger.debug( + f"Atomic check+increment response: {json.dumps(atomic_response, indent=2)}" + ) - # Model-wide limit exceeded (ALWAYS enforce) - if descriptor_key == "model_saturation_check": - raise HTTPException( - status_code=429, - detail={ - "error": f"Model capacity reached for {model}. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - }, - ) + if atomic_response["overall_code"] == "OVER_LIMIT": + for status in atomic_response["statuses"]: + if status["code"] != "OVER_LIMIT": + continue + descriptor_key = status["descriptor_key"] + if descriptor_key == "model_saturation_check": + raise HTTPException( + status_code=429, + detail={ + "error": f"Model capacity reached for {model}. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + }, + ) + if descriptor_key == "priority_model": + verbose_proxy_logger.debug( + f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " + f"priority: {priority}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": f"Priority-based rate limit exceeded. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + ) - # Priority limit exceeded (ONLY enforce when saturated) - elif ( - descriptor_key == "priority_model" - and should_enforce_priority - ): - verbose_proxy_logger.debug( - f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " - f"priority: {priority}" - ) - raise HTTPException( - status_code=429, - detail={ - "error": f"Priority-based rate limit exceeded. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}, " - f"Model saturation: {saturation:.1%}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - "x-litellm-saturation": f"{saturation:.2%}", - }, - ) - - # PHASE 3: Increment counters separately to avoid early-exit issues - # Model counter must ALWAYS increment, but priority counter might be over limit - # If we increment them together, v3_limiter's in-memory check will exit early - # and skip incrementing the model counter - - # Step 3a: Increment model-wide counter (always) - model_increment_response = await self.v3_limiter.should_rate_limit( - descriptors=[model_wide_descriptor], + # If priority is NOT enforced (saturation below threshold) but + # priority_descriptors exist, increment them for tracking only — no + # check, no rollback. This matches the prior tracking semantics. + if priority_descriptors and not should_enforce_priority: + priority_tracking_response = await self.v3_limiter.should_rate_limit( + descriptors=priority_descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=False, ) - - # Step 3b: Increment priority counter (may be over limit, but we still track it) - if priority_descriptors: - priority_increment_response = await self.v3_limiter.should_rate_limit( - descriptors=priority_descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=False, - ) - - # Combine responses for post-call hook - combined_response = { - "overall_code": model_increment_response["overall_code"], - "statuses": model_increment_response["statuses"] - + priority_increment_response["statuses"], - } - data["litellm_proxy_rate_limit_response"] = combined_response - else: - data["litellm_proxy_rate_limit_response"] = model_increment_response + data["litellm_proxy_rate_limit_response"] = { + "overall_code": atomic_response["overall_code"], + "statuses": atomic_response["statuses"] + + priority_tracking_response["statuses"], + } + else: + data["litellm_proxy_rate_limit_response"] = atomic_response async def async_pre_call_hook( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5fea31608f..9de14f55c9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -81,6 +81,85 @@ end return results """ +CHECK_AND_INCREMENT_BY_N_SCRIPT = """ +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: +-- ARGV[1] = now (unix seconds) +-- ARGV[2] = window_size (seconds) +-- For each descriptor i (1..N), starting at ARGV[3]: +-- ARGV[3 + (i-1)*3 + 0] = limit +-- ARGV[3 + (i-1)*3 + 1] = increment +-- ARGV[3 + (i-1)*3 + 2] = ttl (counter TTL when window resets) +-- +-- Return on success: { 0, new_counter_1, new_counter_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) +local descriptor_count = #KEYS / 2 + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = 3 + (i - 1) * 3 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + if current_counter + increment > limit then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = 3 + (i - 1) * 3 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + + local window_expired = descriptor_state[i][1] + + if window_expired then + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end +end + +return results +""" + TOKEN_INCREMENT_SCRIPT = """ local results = {} @@ -163,9 +242,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TOKEN_INCREMENT_SCRIPT ) ) + self.check_and_increment_by_n_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + CHECK_AND_INCREMENT_BY_N_SCRIPT + ) + ) else: self.batch_rate_limiter_script = None self.token_increment_script = None + self.check_and_increment_by_n_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) @@ -595,6 +680,233 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return rate_limit_response + async def atomic_check_and_increment_by_n( + self, + descriptors: List[RateLimitDescriptor], + increments: List[Dict[Literal["requests", "tokens"], int]], + parent_otel_span: Optional[Span] = None, + ) -> RateLimitResponse: + """ + Atomic check-and-increment-by-N across one or more descriptors. + + All-or-nothing: if any descriptor would exceed its limit, no counter is + modified and the response carries `overall_code = "OVER_LIMIT"` with + the offending descriptor's status. Closes the TOCTOU window between + read and increment in both single-process and multi-process (Redis) + deployments. + + Args: + descriptors: rate-limit descriptors to check + increments: per-descriptor increment amounts, indexed parallel to + `descriptors`. Each entry is `{"requests": int, "tokens": int}` + — values default to 0 when a descriptor has no matching limit. + + Returns: + RateLimitResponse with one status per (descriptor, rate_limit_type) + counter, mirroring `should_rate_limit`'s shape. + """ + if len(descriptors) != len(increments): + raise ValueError( + "atomic_check_and_increment_by_n: descriptors and increments " + "must have the same length" + ) + + keys: List[str] = [] + per_counter_meta: List[Dict[str, Any]] = [] + script_args: List[Any] = [] + + for descriptor, increment_amounts in zip(descriptors, increments): + descriptor_key = descriptor["key"] + descriptor_value = descriptor["value"] + rate_limit: RateLimitDescriptorRateLimitObject = ( + descriptor.get("rate_limit") or RateLimitDescriptorRateLimitObject() + ) + window_size = rate_limit.get("window_size") or self.window_size + window_key = f"{{{descriptor_key}:{descriptor_value}}}:window" + + for rate_limit_type in ("requests", "tokens"): + rlt: Literal["requests", "tokens"] = cast( + Literal["requests", "tokens"], rate_limit_type + ) + if rlt == "requests": + limit_value = rate_limit.get("requests_per_unit") + inc_amount = int(increment_amounts.get("requests", 0) or 0) + else: + limit_value = rate_limit.get("tokens_per_unit") + inc_amount = int(increment_amounts.get("tokens", 0) or 0) + if limit_value is None or inc_amount <= 0: + continue + counter_key = self.create_rate_limit_keys( + descriptor_key, descriptor_value, rlt + ) + keys.extend([window_key, counter_key]) + script_args.extend([int(limit_value), inc_amount, int(window_size)]) + per_counter_meta.append( + { + "descriptor_key": descriptor_key, + "current_limit": int(limit_value), + "rate_limit_type": rlt, + "window_key": window_key, + "counter_key": counter_key, + "increment": inc_amount, + "ttl": int(window_size), + } + ) + + if not keys: + return RateLimitResponse(overall_code="OK", statuses=[]) + + current_time = self._get_current_time() + now_int = int(current_time.timestamp()) + + # Multi-process atomicity via Redis Lua. Single-process atomicity + # falls back to the asyncio.Lock + in-memory sliding window below. + if self.check_and_increment_by_n_script is not None: + try: + raw = await self.check_and_increment_by_n_script( + keys=keys, + args=[now_int, self.window_size] + script_args, + ) + return self._build_atomic_response(raw, per_counter_meta) + except Exception as e: + verbose_proxy_logger.warning( + f"atomic_check_and_increment_by_n Lua failed, falling back " + f"to in-memory: {str(e)}" + ) + + async with self._check_and_increment_lock: + return await self._atomic_check_and_increment_in_memory( + per_counter_meta=per_counter_meta, + now_int=now_int, + parent_otel_span=parent_otel_span, + ) + + def _build_atomic_response( + self, + raw: List[Any], + per_counter_meta: List[Dict[str, Any]], + ) -> RateLimitResponse: + """Convert Lua script return value to RateLimitResponse.""" + if not raw: + return RateLimitResponse(overall_code="OK", statuses=[]) + + status_code = int(raw[0]) + if status_code == 1: + # Over limit: { 1, descriptor_index (1-based), current_counter, limit } + descriptor_index = int(raw[1]) - 1 + current_counter = int(raw[2]) + limit = int(raw[3]) + meta = per_counter_meta[descriptor_index] + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ + RateLimitStatus( + code="OVER_LIMIT", + current_limit=limit, + limit_remaining=max(0, limit - current_counter), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ], + ) + + statuses: List[RateLimitStatus] = [] + for meta, new_counter in zip(per_counter_meta, raw[1:]): + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - int(new_counter)), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + return RateLimitResponse(overall_code="OK", statuses=statuses) + + async def _atomic_check_and_increment_in_memory( + self, + per_counter_meta: List[Dict[str, Any]], + now_int: int, + parent_otel_span: Optional[Span] = None, + ) -> RateLimitResponse: + """In-memory all-or-nothing check-and-increment. Caller holds lock.""" + # Pass 1: read state, validate. + descriptor_state: List[Dict[str, Any]] = [] + for meta in per_counter_meta: + window_start = await self.internal_usage_cache.async_get_cache( + key=meta["window_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + window_expired = ( + window_start is None + or (now_int - int(window_start)) >= self.window_size + ) + current_counter = ( + 0 + if window_expired + else int( + await self.internal_usage_cache.async_get_cache( + key=meta["counter_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + or 0 + ) + ) + if current_counter + meta["increment"] > meta["current_limit"]: + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ + RateLimitStatus( + code="OVER_LIMIT", + current_limit=meta["current_limit"], + limit_remaining=max( + 0, meta["current_limit"] - current_counter + ), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ], + ) + descriptor_state.append( + {"window_expired": window_expired, "current": current_counter} + ) + + # Pass 2: apply increments. + statuses: List[RateLimitStatus] = [] + for meta, state in zip(per_counter_meta, descriptor_state): + new_counter = ( + meta["increment"] + if state["window_expired"] + else state["current"] + meta["increment"] + ) + if state["window_expired"]: + await self.internal_usage_cache.async_set_cache( + key=meta["window_key"], + value=str(now_int), + ttl=self.window_size, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=meta["counter_key"], + value=new_counter, + ttl=meta["ttl"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - new_counter), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + return RateLimitResponse(overall_code="OK", statuses=statuses) + def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: Optional[str] = None ) -> List[RateLimitDescriptor]: diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py index 006dac0449..076cb9d3f1 100644 --- a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -19,7 +19,7 @@ check-and-increment becomes atomic. import asyncio import os import sys -from typing import Any, Dict, List, Optional +from typing import List import pytest @@ -102,9 +102,7 @@ async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou(): tpm_limit=TPM_LIMIT, rpm_limit=1000, ) - batch_usage = BatchFileUsage( - total_tokens=BATCH_TOKENS, request_count=1 - ) + batch_usage = BatchFileUsage(total_tokens=BATCH_TOKENS, request_count=1) barrier = _make_phase1_barrier(NUM_CONCURRENT) rate_limiter.should_rate_limit = barrier(rate_limiter.should_rate_limit) @@ -136,15 +134,13 @@ async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou(): @pytest.mark.asyncio -async def test_batch_limiter_check_and_increment_is_two_separate_calls(): +async def test_batch_limiter_uses_atomic_check_and_increment(): """ - Structural test: _check_and_increment_batch_counters issues a read_only=True - check followed by a separate increment call — non-atomic by construction. + Regression test: batch limiter routes through + `atomic_check_and_increment_by_n` rather than the legacy two-phase + pattern (read_only=True check + separate async_increment_tokens_with_ttl_preservation). - Records call ordering on parallel_request_limiter to prove Phase 1 (check) - and Phase 3 (increment) are not wrapped in a single Redis transaction / - Lua script. After the fix, this pattern should be replaced with one - atomic_check_and_increment call. + Ensures future refactors don't reintroduce the TOCTOU window. """ dual_cache = DualCache() internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) @@ -154,25 +150,23 @@ async def test_batch_limiter_check_and_increment_is_two_separate_calls(): batch_limiter = rate_limiter._get_batch_rate_limiter() assert batch_limiter is not None - call_log: List[Dict[str, Any]] = [] + call_log: List[str] = [] + original_atomic = rate_limiter.atomic_check_and_increment_by_n original_should = rate_limiter.should_rate_limit - original_inc = rate_limiter.async_increment_tokens_with_ttl_preservation + + async def logging_atomic(*args, **kwargs): + call_log.append("atomic_check_and_increment_by_n") + return await original_atomic(*args, **kwargs) async def logging_should(*args, **kwargs): - call_log.append( - {"method": "should_rate_limit", "read_only": kwargs.get("read_only")} - ) + call_log.append(f"should_rate_limit(read_only={kwargs.get('read_only')})") return await original_should(*args, **kwargs) - async def logging_inc(*args, **kwargs): - call_log.append({"method": "async_increment_tokens_with_ttl_preservation"}) - return await original_inc(*args, **kwargs) - + rate_limiter.atomic_check_and_increment_by_n = logging_atomic rate_limiter.should_rate_limit = logging_should - rate_limiter.async_increment_tokens_with_ttl_preservation = logging_inc user_api_key_dict = UserAPIKeyAuth( - api_key=hash_token("structural-test-key"), + api_key=hash_token("atomic-test-key"), tpm_limit=10000, rpm_limit=1000, ) @@ -183,29 +177,14 @@ async def test_batch_limiter_check_and_increment_is_two_separate_calls(): batch_usage=BatchFileUsage(total_tokens=50, request_count=1), ) - method_sequence = [c["method"] for c in call_log] - assert "should_rate_limit" in method_sequence, "Expected Phase 1 check" - phase1 = [ - c - for c in call_log - if c["method"] == "should_rate_limit" and c.get("read_only") is True - ] - phase3 = [ - c - for c in call_log - if c["method"] == "async_increment_tokens_with_ttl_preservation" - ] - assert len(phase1) >= 1 and len(phase3) >= 1, ( - f"Expected non-atomic Phase1+Phase3 pattern. call_log={call_log}" + assert "atomic_check_and_increment_by_n" in call_log, ( + f"Batch limiter must route through atomic_check_and_increment_by_n. " + f"Calls observed: {call_log}" ) - phase1_idx = method_sequence.index("should_rate_limit") - phase3_idx = method_sequence.index( - "async_increment_tokens_with_ttl_preservation" - ) - assert phase1_idx < phase3_idx, ( - "TOCTOU evidence: read-only check precedes increment as separate awaits — " - "no atomic Lua script wraps both. Sequence: " - f"{method_sequence}" + legacy_calls = [c for c in call_log if c.startswith("should_rate_limit(")] + assert not legacy_calls, ( + f"Batch limiter must not call should_rate_limit directly (legacy " + f"two-phase pattern). Observed: {legacy_calls}" ) @@ -295,15 +274,15 @@ async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_phase1_phase3_are_separate_awaits(): +async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): """ - Structural proof of TOCTOU: dynamic_rate_limiter_v3._check_rate_limits - issues a read_only=True call (Phase 1, line 464-468) followed by separate - read_only=False calls (Phase 3, lines 526-530 + 534-538). + Regression test: dynamic limiter's enforced descriptors flow through + `atomic_check_and_increment_by_n`, not the legacy + read_only=True check followed by a separate read_only=False increment. - Records each invocation of v3_limiter.should_rate_limit and asserts the - Phase1→Phase3 sequence. After fix, both phases must collapse into a single - atomic operation. + When priority is enforced (saturation >= threshold), priority_model is + bundled into the atomic call alongside model_saturation_check. When not + enforced, priority counter is incremented for tracking only. """ os.environ["LITELLM_LICENSE"] = "test-license-key" litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -311,7 +290,7 @@ async def test_dynamic_rate_limiter_v3_phase1_phase3_are_separate_awaits(): dual_cache = DualCache() handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) - model = "structural-dyn-model" + model = "atomic-dyn-model" llm_router = Router( model_list=[ { @@ -327,18 +306,19 @@ async def test_dynamic_rate_limiter_v3_phase1_phase3_are_separate_awaits(): ) handler.update_variables(llm_router=llm_router) - read_only_flags: List[Optional[bool]] = [] - original = handler.v3_limiter.should_rate_limit + atomic_descriptors_observed: List[List[str]] = [] + original_atomic = handler.v3_limiter.atomic_check_and_increment_by_n - async def logging_should(*args, **kwargs): - read_only_flags.append(kwargs.get("read_only")) - return await original(*args, **kwargs) + async def logging_atomic(*args, **kwargs): + ds = kwargs.get("descriptors") or (args[0] if args else []) + atomic_descriptors_observed.append([d["key"] for d in ds]) + return await original_atomic(*args, **kwargs) - handler.v3_limiter.should_rate_limit = logging_should + handler.v3_limiter.atomic_check_and_increment_by_n = logging_atomic from litellm.types.router import ModelGroupInfo - user = UserAPIKeyAuth(api_key=hash_token("dyn-structural-key")) + user = UserAPIKeyAuth(api_key=hash_token("dyn-atomic-key")) user.metadata = {"priority": "high"} await handler._check_rate_limits( @@ -355,21 +335,12 @@ async def test_dynamic_rate_limiter_v3_phase1_phase3_are_separate_awaits(): data={}, ) - assert True in read_only_flags or any( - f is True for f in read_only_flags - ), f"Expected read_only=True (Phase 1) call. Got: {read_only_flags}" - assert any(f is False for f in read_only_flags), ( - f"Expected read_only=False (Phase 3) call. Got: {read_only_flags}" + assert atomic_descriptors_observed, ( + "Dynamic limiter must route enforced descriptors through " + "atomic_check_and_increment_by_n (no legacy read_only=True / " + "separate-increment pattern)." ) - - first_read_only = next( - (i for i, f in enumerate(read_only_flags) if f is True), None - ) - first_write = next( - (i for i, f in enumerate(read_only_flags) if f is False), None - ) - assert first_read_only is not None and first_write is not None - assert first_read_only < first_write, ( - f"TOCTOU evidence: Phase 1 (read_only) precedes Phase 3 (increment) " - f"as separate non-atomic calls. read_only sequence: {read_only_flags}" + assert "model_saturation_check" in atomic_descriptors_observed[0], ( + f"Expected model_saturation_check in atomic descriptor set. " + f"Got: {atomic_descriptors_observed}" ) From 6496e584175b4ec1289443a02feaef3872915ee0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 1 May 2026 12:02:04 -0700 Subject: [PATCH 3/5] review: address atomic limiter review feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Lua script now reads time via redis.call('TIME') instead of a client-supplied timestamp. Prevents window-reset divergence across replicas with skewed wall-clocks, which could otherwise reopen the cross-replica TOCTOU window. - Per-descriptor window_size is now plumbed through both the Lua ARGV layout and the in-memory fallback. Previously the in-memory path used the global self.window_size while Lua honored the per-descriptor override, so a descriptor with a custom window would be enforced inconsistently between Redis-available and Redis-unavailable code paths. - Lua-failure fallback path now logs at error severity and explicitly documents the in-memory ↔ Redis state divergence risk so operators can alert on it. Prior `warning` log understated the impact. - Coarse-granularity lock is now documented inline with the conditions under which a per-descriptor sharded lock would be worth introducing. - New regression test: zero-token batch consumes RPM only and is properly capped by the RPM ceiling (validates the asymmetric quota path that arises from `inc_amount <= 0: continue`). Tests: 64 passed (1 skipped), 0 regressions. Multi-instance Redis loadtest re-verified: chat 20/80 success @ RPM=20, batches 3/20 @ TPM=200. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../hooks/parallel_request_limiter_v3.py | 95 ++++++++++++++----- .../proxy/hooks/test_rate_limiter_toctou.py | 68 +++++++++++++ 2 files changed, 138 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9de14f55c9..cd85843ef1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -86,19 +86,22 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT = """ -- All-or-nothing: if any descriptor would exceed its limit, no counter is -- modified. -- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- -- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. --- ARGV layout: --- ARGV[1] = now (unix seconds) --- ARGV[2] = window_size (seconds) --- For each descriptor i (1..N), starting at ARGV[3]: --- ARGV[3 + (i-1)*3 + 0] = limit --- ARGV[3 + (i-1)*3 + 1] = increment --- ARGV[3 + (i-1)*3 + 2] = ttl (counter TTL when window resets) +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) -- -- Return on success: { 0, new_counter_1, new_counter_2, ... } -- Return on over-limit: { 1, descriptor_index, current_counter, limit } -local now = tonumber(ARGV[1]) -local window_size = tonumber(ARGV[2]) +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) local descriptor_count = #KEYS / 2 -- Pass 1: read state, validate. Abort without writing if any over limit. @@ -106,9 +109,10 @@ local descriptor_state = {} for i = 1, descriptor_count do local window_key = KEYS[(i - 1) * 2 + 1] local counter_key = KEYS[(i - 1) * 2 + 2] - local arg_base = 3 + (i - 1) * 3 + local arg_base = (i - 1) * 4 + 1 local limit = tonumber(ARGV[arg_base]) local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) local window_start = redis.call('GET', window_key) local window_expired = (not window_start) or @@ -133,9 +137,10 @@ local results = { 0 } for i = 1, descriptor_count do local window_key = KEYS[(i - 1) * 2 + 1] local counter_key = KEYS[(i - 1) * 2 + 2] - local arg_base = 3 + (i - 1) * 3 + local arg_base = (i - 1) * 4 + 1 local increment = tonumber(ARGV[arg_base + 1]) local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) local window_expired = descriptor_state[i][1] @@ -261,6 +266,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # limiters) within this process to close the TOCTOU window between # read-only check and counter increment. Multi-replica deployments # additionally rely on Redis Lua atomicity for cross-process safety. + # + # Coarse granularity: this single lock serializes ALL atomic check+ + # increment operations across batch and dynamic limiters on this + # instance. A slow batch input-file fetch (which happens upstream of + # the lock) does not block here, but Redis Lua latency does. If + # contention shows up under load (visible as p99 latency spikes + # correlated with batch traffic), shard to a per-descriptor-key lock + # via a `weakref.WeakValueDictionary[str, asyncio.Lock]`. Punted as a + # follow-up because Lua dominates wall-time and the lock is held for + # one round-trip. self._check_and_increment_lock = asyncio.Lock() def _get_batch_rate_limiter(self) -> Optional[Any]: @@ -740,7 +755,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptor_key, descriptor_value, rlt ) keys.extend([window_key, counter_key]) - script_args.extend([int(limit_value), inc_amount, int(window_size)]) + # Per-descriptor 4-tuple: limit, increment, ttl, window_size. + # window_size is per-descriptor — descriptors may carry custom + # windows distinct from self.window_size, and the Lua script + # uses this slot to evaluate window expiry. + script_args.extend( + [ + int(limit_value), + inc_amount, + int(window_size), + int(window_size), + ] + ) per_counter_meta.append( { "descriptor_key": descriptor_key, @@ -750,34 +776,44 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "counter_key": counter_key, "increment": inc_amount, "ttl": int(window_size), + "window_size": int(window_size), } ) if not keys: return RateLimitResponse(overall_code="OK", statuses=[]) - current_time = self._get_current_time() - now_int = int(current_time.timestamp()) - # Multi-process atomicity via Redis Lua. Single-process atomicity # falls back to the asyncio.Lock + in-memory sliding window below. + # Note: in-memory state diverges from Redis state — if Lua fails + # mid-write, retrying via in-memory may double-count. See fallback + # warning below. if self.check_and_increment_by_n_script is not None: try: raw = await self.check_and_increment_by_n_script( keys=keys, - args=[now_int, self.window_size] + script_args, + args=script_args, ) return self._build_atomic_response(raw, per_counter_meta) except Exception as e: - verbose_proxy_logger.warning( - f"atomic_check_and_increment_by_n Lua failed, falling back " - f"to in-memory: {str(e)}" + # Escalated from warning to error: Lua failures (script timeout, + # Redis OOM, network partition) leave counter state ambiguous. + # The fallback path below uses LOCAL DualCache, which is a + # different store from Redis — counters here will diverge from + # Redis until that key's window expires (TTL bounds divergence). + # Operators should alert on this log line; sustained occurrences + # indicate Redis health degradation that may erode rate-limit + # accuracy. + verbose_proxy_logger.error( + f"atomic_check_and_increment_by_n: Redis Lua execution " + f"failed ({type(e).__name__}: {e}). Falling back to " + f"in-memory enforcement — counters will diverge from Redis " + f"state until window expires (window_size={self.window_size}s)." ) async with self._check_and_increment_lock: return await self._atomic_check_and_increment_in_memory( per_counter_meta=per_counter_meta, - now_int=now_int, parent_otel_span=parent_otel_span, ) @@ -826,21 +862,30 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _atomic_check_and_increment_in_memory( self, per_counter_meta: List[Dict[str, Any]], - now_int: int, parent_otel_span: Optional[Span] = None, ) -> RateLimitResponse: - """In-memory all-or-nothing check-and-increment. Caller holds lock.""" + """In-memory all-or-nothing check-and-increment. Caller holds lock. + + Reads/writes the LOCAL DualCache (`local_only=True`) — note this is + a different store from Redis. When this fallback fires after a Lua + failure, in-memory counters will diverge from Redis until each key's + window expires (TTL bounds divergence). + """ + # Use a single 'now' for the duration of this critical section so all + # descriptors evaluate window expiry consistently. + now_int = int(self._get_current_time().timestamp()) + # Pass 1: read state, validate. descriptor_state: List[Dict[str, Any]] = [] for meta in per_counter_meta: + window_size = meta["window_size"] window_start = await self.internal_usage_cache.async_get_cache( key=meta["window_key"], litellm_parent_otel_span=parent_otel_span, local_only=True, ) window_expired = ( - window_start is None - or (now_int - int(window_start)) >= self.window_size + window_start is None or (now_int - int(window_start)) >= window_size ) current_counter = ( 0 @@ -885,7 +930,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): await self.internal_usage_cache.async_set_cache( key=meta["window_key"], value=str(now_int), - ttl=self.window_size, + ttl=meta["window_size"], litellm_parent_otel_span=parent_otel_span, local_only=True, ) diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py index 076cb9d3f1..b8ced4e661 100644 --- a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -344,3 +344,71 @@ async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): f"Expected model_saturation_check in atomic descriptor set. " f"Got: {atomic_descriptors_observed}" ) + + +@pytest.mark.asyncio +async def test_batch_zero_token_consumes_rpm_only(): + """ + Zero-token batch (e.g. metadata-only call) should still increment RPM + counter but NOT TPM counter. + + Edge case from review: `if inc_amount <= 0: continue` in + `atomic_check_and_increment_by_n` skips descriptor counters whose + increment is zero. Verifies asymmetric quota consumption is intentional + and observable: an RPM-bounded but TPM-free request path stays bounded + by RPM alone. + """ + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("zero-token-key"), + tpm_limit=100, + rpm_limit=3, + ) + zero_batch = BatchFileUsage(total_tokens=0, request_count=1) + + # 3 zero-token batches must succeed (RPM=3 allows). 4th must hit RPM cap, + # NOT TPM (because token counter never increments past 0). + for i in range(3): + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=zero_batch, + ) + + # Inspect counters: RPM key incremented to 3, TPM key absent (or 0). + rpm_key = rate_limiter.create_rate_limit_keys( + "api_key", user_api_key_dict.api_key or "", "requests" + ) + tpm_key = rate_limiter.create_rate_limit_keys( + "api_key", user_api_key_dict.api_key or "", "tokens" + ) + rpm_val = await internal_usage_cache.async_get_cache( + key=rpm_key, litellm_parent_otel_span=None, local_only=True + ) + tpm_val = await internal_usage_cache.async_get_cache( + key=tpm_key, litellm_parent_otel_span=None, local_only=True + ) + assert int(rpm_val or 0) == 3, f"RPM counter must reach 3, got {rpm_val}" + assert tpm_val in ( + None, + 0, + "0", + ), f"TPM counter must remain unset/0 for zero-token batches, got {tpm_val}" + + # 4th attempt: RPM exhausted -> 429. + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=zero_batch, + ) + assert exc.value.status_code == 429 From eba0cdf3f5e091c9b656a05cc9a9e9082f15fa6d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 1 May 2026 12:19:43 -0700 Subject: [PATCH 4/5] fix(rate-limit): fail closed on unrecognized OVER_LIMIT descriptor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit If atomic_check_and_increment_by_n returns overall_code=OVER_LIMIT but no status entry matches a descriptor key the dynamic limiter dispatcher knows how to translate into a 429 (`model_saturation_check` or `priority_model`), the for-loop previously exited cleanly and execution fell through to the priority-tracking increment + the data["litellm_proxy_rate_limit_response"] write — silently admitting an over-limit request. This is the fail-open path a future contributor would hit by wiring a new descriptor type into enforced_descriptors without updating the dispatcher. Refuse the request with a generic 429 carrying the offending descriptor metadata so the operator can see what slipped past, and emit an error log to surface the wiring gap. Adds a regression test (test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor) that drives the limiter with a synthetic OVER_LIMIT response carrying an unrecognized descriptor_key and asserts a 429 is raised. Tests: 65 passed (1 skipped), 0 regressions. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 31 ++++++++ .../proxy/hooks/test_rate_limiter_toctou.py | 75 +++++++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index d827cf3067..2cbb527284 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -528,6 +528,37 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): }, ) + # Fail-closed guard: overall_code says OVER_LIMIT but no status + # matched a descriptor key we know how to translate into a 429. + # Refuse the request rather than silently fall through and let an + # over-limit request proceed to the model. Without this, a future + # caller wiring an unfamiliar descriptor into enforced_descriptors + # would silently bypass the rate limit. + offending = next( + (s for s in atomic_response["statuses"] if s["code"] == "OVER_LIMIT"), + None, + ) + verbose_proxy_logger.error( + f"Dynamic rate limiter: OVER_LIMIT response with unknown " + f"descriptor_key(s) — refusing request. response={atomic_response}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": "Rate limit exceeded", + "descriptor_key": ( + offending["descriptor_key"] if offending else "unknown" + ), + "rate_limit_type": ( + str(offending["rate_limit_type"]) if offending else "unknown" + ), + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "x-litellm-priority": priority or "default", + }, + ) + # If priority is NOT enforced (saturation below threshold) but # priority_descriptors exist, increment them for tracking only — no # check, no rollback. This matches the prior tracking semantics. diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py index b8ced4e661..ceea5de799 100644 --- a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -412,3 +412,78 @@ async def test_batch_zero_token_consumes_rpm_only(): batch_usage=zero_batch, ) assert exc.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): + """ + Fail-closed guard: when atomic_check_and_increment_by_n returns + overall_code=OVER_LIMIT but with a descriptor_key the dispatcher does + not recognize, the dynamic limiter must raise 429 rather than silently + fall through. + + Reproduces by patching atomic_check_and_increment_by_n to return an + OVER_LIMIT response carrying an unknown descriptor_key. + """ + from fastapi import HTTPException + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fail-closed-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + async def fake_atomic(*args, **kwargs): + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "tokens", + "descriptor_key": "future_unrecognized_descriptor", + } + ], + } + + handler.v3_limiter.atomic_check_and_increment_by_n = fake_atomic + + from litellm.types.router import ModelGroupInfo + + user = UserAPIKeyAuth(api_key=hash_token("fail-closed-key")) + user.metadata = {"priority": "high"} + + with pytest.raises(HTTPException) as exc: + await handler._check_rate_limits( + model=model, + model_group_info=ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=None, + tpm=1000, + ), + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + assert ( + exc.value.status_code == 429 + ), f"Expected 429 fail-closed on unknown descriptor; got {exc.value.status_code}" From ca3e659a3cf41f0e2e9817334d33605ee88486f9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 1 May 2026 12:40:43 -0700 Subject: [PATCH 5/5] review: separate ttl_seconds from window_size_seconds in atomic ARGV MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Lua ARGV layout for `CHECK_AND_INCREMENT_BY_N_SCRIPT` documents two distinct slots — `ttl_seconds` (counter-key TTL) and `window_size_seconds` (sliding-window length) — but the Python call site collapsed both to `int(window_size)`. They happen to be equal today, so the bug was latent, but a future descriptor carrying a custom counter TTL would have been silently mis-applied (counter key expiring at the wrong time relative to the window). Make the two values separate variables (`ttl_seconds`, `window_size_seconds`) before the `script_args.extend([...])` call so the ARGV layout matches the documented contract by construction. The values still flow through to `per_counter_meta` so the in-memory fallback uses the same source of truth. Also: - Document the COUNTER-level (not descriptor-level) parallel-indexing invariant on `_build_atomic_response` and `per_counter_meta`. A descriptor with both RPM and TPM emits two counters and two meta entries; the Lua script's `i` enumerates counters and the over-limit return tuple's index maps directly to `per_counter_meta[i - 1]`. Comment surfaces this so a future contributor doesn't refactor the arrays apart. - Annotate the `increments` literal in the batch limiter directly instead of routing through `cast(...)`. Drops an unneeded `cast` import. - Comment on the dynamic limiter's tracking-only branch explaining why `should_rate_limit` (non-atomic) is intentional there: we want to bump the priority counter so future enforced windows have accurate usage, and the OVER_LIMIT response is intentionally ignored. Tests: 65 passed (1 skipped), 0 regressions. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/hooks/batch_rate_limiter.py | 19 ++++------ .../proxy/hooks/dynamic_rate_limiter_v3.py | 7 ++++ .../hooks/parallel_request_limiter_v3.py | 38 ++++++++++++++----- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index ffbcfc0f44..2ee0588f19 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,7 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union from fastapi import HTTPException from pydantic import BaseModel @@ -179,16 +179,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): model_has_failures=False, ) - increments = cast( - List[Dict[Literal["requests", "tokens"], int]], - [ - { - "requests": batch_usage.request_count, - "tokens": batch_usage.total_tokens, - } - for _ in descriptors - ], - ) + increment: Dict[Literal["requests", "tokens"], int] = { + "requests": batch_usage.request_count, + "tokens": batch_usage.total_tokens, + } + increments: List[Dict[Literal["requests", "tokens"], int]] = [ + increment for _ in descriptors + ] rate_limit_response = ( await self.parallel_request_limiter.atomic_check_and_increment_by_n( diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 2cbb527284..f7c0592992 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -562,6 +562,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # If priority is NOT enforced (saturation below threshold) but # priority_descriptors exist, increment them for tracking only — no # check, no rollback. This matches the prior tracking semantics. + # + # Using the non-atomic should_rate_limit (instead of + # atomic_check_and_increment_by_n) is intentional here: we don't want + # to enforce the limit, we only want to bump the counter so the + # priority allocation has accurate usage when it later becomes + # enforced. The increment-then-check semantics of should_rate_limit + # are fine because we ignore the OVER_LIMIT response. if priority_descriptors and not should_enforce_priority: priority_tracking_response = await self.v3_limiter.should_rate_limit( descriptors=priority_descriptors, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index cd85843ef1..4497e64c17 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -754,17 +754,25 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): counter_key = self.create_rate_limit_keys( descriptor_key, descriptor_value, rlt ) + # Counter-key TTL and window_size are conceptually distinct + # ("how long the counter Redis key lives" vs "how long the + # sliding window is"). They happen to be equal today because + # we have no descriptor type that needs them apart, but they + # are kept as separate variables here so a future custom-TTL + # descriptor doesn't reintroduce a silent expiry bug. Both + # the Lua script and the in-memory fallback read these from + # their respective ARGV / meta slots. + ttl_seconds = int(window_size) + window_size_seconds = int(window_size) keys.extend([window_key, counter_key]) - # Per-descriptor 4-tuple: limit, increment, ttl, window_size. - # window_size is per-descriptor — descriptors may carry custom - # windows distinct from self.window_size, and the Lua script - # uses this slot to evaluate window expiry. + # Per-counter 4-tuple matches the Lua ARGV layout exactly: + # [limit, increment, ttl_seconds, window_size_seconds]. script_args.extend( [ int(limit_value), inc_amount, - int(window_size), - int(window_size), + ttl_seconds, + window_size_seconds, ] ) per_counter_meta.append( @@ -775,8 +783,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "window_key": window_key, "counter_key": counter_key, "increment": inc_amount, - "ttl": int(window_size), - "window_size": int(window_size), + "ttl": ttl_seconds, + "window_size": window_size_seconds, } ) @@ -822,13 +830,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): raw: List[Any], per_counter_meta: List[Dict[str, Any]], ) -> RateLimitResponse: - """Convert Lua script return value to RateLimitResponse.""" + """Convert Lua script return value to RateLimitResponse. + + Indexing invariant: `per_counter_meta` and `KEYS` are parallel-indexed + at the COUNTER level, not the descriptor level. A descriptor with both + RPM and TPM limits emits two `(window_key, counter_key)` pairs and + two meta entries — one per counter. The Lua script's loop variable + `i` therefore enumerates counters, and the over-limit return tuple + `{1, i, ...}` carries a counter index that maps directly to + `per_counter_meta[i - 1]`. Keep these arrays parallel at the counter + level when modifying this code. + """ if not raw: return RateLimitResponse(overall_code="OK", statuses=[]) status_code = int(raw[0]) if status_code == 1: - # Over limit: { 1, descriptor_index (1-based), current_counter, limit } + # Over limit: { 1, counter_index (1-based), current_counter, limit } descriptor_index = int(raw[1]) - 1 current_counter = int(raw[2]) limit = int(raw[3])