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