From d8ea14d8ff96519fb158fce6a7c0983d135fff36 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 20 Aug 2026 16:42:15 -0400 Subject: [PATCH] fix(rate-limiting): read token/dollar buckets straight from Redis, not a stale in-memory snapshot Success accounting increments these buckets through a Lua script that writes directly to Redis, bypassing DualCache's in-memory layer entirely. Once an earlier admission read had backfilled that key into the in-memory cache, DualCache's own batch-read treats that non-None hit as authoritative and never rechecks Redis, so later admissions kept seeing the same frozen snapshot while the real counter climbed underneath it, silently admitting traffic past the configured token/dollar limit for up to the in-memory TTL. Admission now reads these buckets directly off the Redis connection when one is configured, skipping the in-memory layer that write path never keeps coherent; the no-Redis, single-process fallback is unaffected since its increments already go through the same in-memory cache these reads use. --- litellm/proxy/hooks/tag_rate_limiter.py | 37 +++++++-- .../proxy/hooks/test_tag_rate_limiter.py | 79 +++++++++++++++++++ 2 files changed, 109 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/tag_rate_limiter.py b/litellm/proxy/hooks/tag_rate_limiter.py index f36b872f5ee..6395db08e24 100644 --- a/litellm/proxy/hooks/tag_rate_limiter.py +++ b/litellm/proxy/hooks/tag_rate_limiter.py @@ -1122,13 +1122,36 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer # LIT010/Final-in-loop check forbids partition = await self._partition_for(partition_key) keys = [read_only_checks[i][2] for i in indices] # mutable-ok: async_batch_get_cache needs a real list - current_values = await partition.internal_usage_cache.async_batch_get_cache( - keys=keys, - parent_otel_span=parent_otel_span, - local_only=False, - ) - missing = [None] * len(keys) # mutable-ok: async_batch_get_cache requires a real list; see above - resolved = current_values if current_values is not None else missing + redis_cache = partition.internal_usage_cache.dual_cache.redis_cache + if redis_cache is not None: + # async_log_success_event increments these buckets straight + # through a Lua script on this same redis_cache, bypassing + # DualCache/InternalUsageCache entirely -- so its in-memory + # layer never learns about that write. DualCache's own + # async_batch_get_cache treats any non-None in-memory hit as + # authoritative and never re-checks Redis for that key (see + # _reserve_redis_batch_keys), so once a key is backfilled + # in-memory it silently freezes for up to the in-memory TTL + # (10 minutes by default) while the real Redis counter keeps + # moving underneath it -- reading straight from Redis here, + # bypassing that in-memory layer, is the only way this + # read-then-later-increment split stays coherent. + # not `Final`: rebound each loop iteration, which basedpyright's + # LIT010/Final-in-loop check forbids; explicitly typed since + # RedisCache.async_batch_get_cache's own signature returns a + # bare, unparameterized dict + redis_values: dict[str, object] = await redis_cache.async_batch_get_cache( + key_list=keys, parent_otel_span=parent_otel_span + ) + resolved = [redis_values.get(key) for key in keys] # mutable-ok: needs a real list + else: + current_values = await partition.internal_usage_cache.async_batch_get_cache( + keys=keys, + parent_otel_span=parent_otel_span, + local_only=True, + ) + missing = [None] * len(keys) # mutable-ok: async_batch_get_cache requires a real list; see above + resolved = current_values if current_values is not None else missing for i, value in zip(indices, resolved): values_by_index[i] = value # mutable-ok: see comment above diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py index 369f0477016..49ca7fb934b 100644 --- a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py @@ -2075,6 +2075,85 @@ async def test_redis_backed_cross_unit_rejection_does_not_leave_a_phantom_increm await redis_cache.async_delete_cache(key=request_key) +@pytest.mark.asyncio +async def test_redis_backed_token_admission_sees_increments_the_in_memory_cache_missed(time_controller): + """ + Success accounting increments a token bucket straight through a Lua + script on redis_cache, bypassing DualCache/InternalUsageCache entirely -- + that write never touches the in-memory layer. Once an earlier read has + backfilled that same key into the in-memory cache, DualCache's own + async_batch_get_cache treats that non-None in-memory hit as authoritative + and never re-checks Redis, so every later admission would see the same + frozen snapshot while the real Redis counter keeps climbing underneath + it, silently admitting traffic well past the configured token limit. + """ + limiter, redis_cache = _redis_limiter(time_controller) + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {e!s}") + + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + {"token_limits": {"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 100, "period_seconds": 60}]}}, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + tag = f"redis-stale-check-{uuid.uuid4().hex}" + request_kwargs = {"metadata": {"tags": [f"end_user_id:{tag}"]}} + + async def _charge(tokens: float) -> None: + await limiter.async_log_success_event( + kwargs={ + "metadata": {"tags": [f"end_user_id:{tag}"]}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": tokens, + "response_cost": 0, + }, + }, + response_obj=None, + start_time=0, + end_time=0, + ) + # The actual Redis increment is dispatched as a background task (see + # _BACKGROUND_TASKS), so it needs a beat to actually run. + await asyncio.sleep(0.05) + + # First admission: bucket doesn't exist in Redis yet, so this read finds + # nothing to backfill into the in-memory cache either. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + await _charge(90) + + # Second admission: this read is the one that backfills the in-memory + # cache with the real (90) value read from Redis. + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + await _charge(90) # real Redis total is now 180, well past the limit of 100 + + # Third admission must see the real (180) total and reject -- not the + # frozen 90 the in-memory cache captured on the previous read. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "per_minute", "end_user_id", tag, 60, now) + await redis_cache.async_delete_cache(key=token_key) + + # --------------------------------------------------------------------------- # team_public_model_name alias -- index lookup must not miss # ---------------------------------------------------------------------------