diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 5b7d7d04e24..1f993cb6bc6 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -11,6 +11,7 @@ Follows the same pattern as parallel_request_limiter_v3.py. """ import os +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final from litellm import DualCache @@ -21,6 +22,8 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.redis.stores.iteration_count import IterationCountStore, TransportIterationCountStore, count_key +from litellm.redis.transport.contract import RedisFailure if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -30,21 +33,6 @@ else: InternalUsageCache = Any -# Redis Lua script for atomic increment with TTL. -# Returns the new count after increment. -# Only sets EXPIRE on first increment (when count becomes 1). -MAX_ITERATIONS_INCREMENT_SCRIPT: Final = """ -local key = KEYS[1] -local ttl = tonumber(ARGV[1]) - -local current = redis.call('INCR', key) -if current == 1 then - redis.call('EXPIRE', key, ttl) -end - -return current -""" - # Default TTL for session iteration counters (1 hour) DEFAULT_MAX_ITERATIONS_TTL: Final = 3600 @@ -64,22 +52,18 @@ class _PROXY_MaxIterationsHandler(CustomLogger): {session_iterations:}:count Multi-instance support: - Uses Redis Lua script for atomic increment (same pattern as - parallel_request_limiter_v3). Falls back to in-memory cache - when Redis is unavailable. + Counts in Redis through the iteration count store, falls back to + in-memory cache when Redis is unavailable. """ def __init__(self, internal_usage_cache: InternalUsageCache): self.internal_usage_cache = internal_usage_cache self.ttl = int(os.getenv("LITELLM_MAX_ITERATIONS_TTL", DEFAULT_MAX_ITERATIONS_TTL)) - # Register Lua script with Redis if available (same pattern as v3 limiter) - if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - MAX_ITERATIONS_INCREMENT_SCRIPT - ) - else: - self.increment_script = None + redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache + self._store: Final[IterationCountStore | None] = ( + TransportIterationCountStore(redis_cache.transport()) if redis_cache is not None else None + ) @with_service_target("session_iterations") async def async_pre_call_hook( @@ -111,8 +95,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): ) # Increment and check - cache_key: Final = self._make_cache_key(session_id) - current_count: Final = await self._increment_and_get(cache_key) + current_count: Final = await self._increment_and_get(session_id) if current_count > max_iterations: resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) @@ -174,36 +157,16 @@ class _PROXY_MaxIterationsHandler(CustomLogger): return None def _make_cache_key(self, session_id: str) -> str: - """ - Create cache key for session iteration counter. + return count_key(session_id) - Uses Redis hash-tag pattern {session_iterations:} so all - keys for a session land on the same Redis Cluster slot. - """ - return f"{{session_iterations:{session_id}}}:count" + async def _increment_and_get(self, session_id: str) -> int: + """Increment the session counter in Redis, or in memory when Redis is unavailable.""" + if self._store is not None: + count: Final = await self._store.increment(session_id, ttl=timedelta(seconds=self.ttl)) + if not isinstance(count, RedisFailure): + return count - async def _increment_and_get(self, cache_key: str) -> int: - """ - Atomically increment the session counter and return the new value. - - Tries Redis first (via registered Lua script for atomicity across - instances), falls back to in-memory cache. - """ - if self.increment_script is not None: - try: - result: Final = await self.increment_script( - keys=[cache_key], - args=[self.ttl], - ) - return int(result) - except Exception as e: - verbose_proxy_logger.warning( - "MaxIterationsHandler: Redis failed, falling back to in-memory: %s", - str(e), - ) - - # Fallback: in-memory cache - return await self._in_memory_increment(cache_key) + return await self._in_memory_increment(self._make_cache_key(session_id)) async def _in_memory_increment(self, cache_key: str) -> int: """Increment counter in in-memory cache with TTL.""" diff --git a/litellm/redis/stores/iteration_count.py b/litellm/redis/stores/iteration_count.py new file mode 100644 index 00000000000..679f01c7c0d --- /dev/null +++ b/litellm/redis/stores/iteration_count.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_INCREMENT: Final = """ +local key = KEYS[1] +local ttl = tonumber(ARGV[1]) + +local current = redis.call('INCR', key) +if current == 1 then + redis.call('EXPIRE', key, ttl) +end + +return current +""" + + +def count_key(session_id: str) -> str: + return f"{{session_iterations:{session_id}}}:count" + + +class IterationCountStore(Protocol): + async def increment(self, session_id: str, *, ttl: timedelta) -> int | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportIterationCountStore: + transport: RedisTransport + + async def increment(self, session_id: str, *, ttl: timedelta) -> int | RedisFailure: + reply: Final = await self.transport.eval(_INCREMENT, (count_key(session_id),), (ttl // timedelta(seconds=1),)) + if isinstance(reply, (RedisFailure, int)): + return reply + return RedisFailure(kind="decode", message=f"session iteration count is not an integer: {reply!r}") diff --git a/tests/unit/proxy/hooks/test_max_iterations_limiter.py b/tests/unit/proxy/hooks/test_max_iterations_limiter.py index 5eb499042b5..661c727afb8 100644 --- a/tests/unit/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/unit/proxy/hooks/test_max_iterations_limiter.py @@ -6,8 +6,9 @@ Tests that session-scoped iteration counting works correctly: - Different sessions have independent counters """ -from unittest.mock import MagicMock, patch +from unittest.mock import patch +import fakeredis import pytest from fastapi import HTTPException @@ -16,6 +17,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandler from litellm.proxy.utils import InternalUsageCache from litellm.types.agents import AgentResponse +from tests._support.fake_redis import fake_redis_cache def _make_mock_agent(max_iterations: int) -> AgentResponse: @@ -150,3 +152,56 @@ async def test_max_iterations_no_agent_id_passes(): call_type="", ) assert result is None + + +def _redis_handler(server: fakeredis.FakeServer) -> PROXY_MaxIterationsHandler: + cache = fake_redis_cache(server, namespace="ns") + return PROXY_MaxIterationsHandler(internal_usage_cache=InternalUsageCache(DualCache(redis_cache=cache))) + + +async def _call(handler: PROXY_MaxIterationsHandler, session_id: str, max_iterations: int) -> None: + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry") as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_iterations=max_iterations) + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", agent_id="agent-test-123"), + cache=DualCache(), + data={"metadata": {"session_id": session_id}}, + call_type="", + ) + + +@pytest.mark.asyncio +async def test_iterations_count_in_redis_with_the_ttl_and_include_an_old_pods_calls(): + server = fakeredis.FakeServer() + raw = fakeredis.FakeAsyncRedis(server=server) + handler = _redis_handler(server) + await raw.incr("ns:{session_iterations:s1}:count") + + await _call(handler, "s1", max_iterations=3) + await _call(handler, "s1", max_iterations=3) + with pytest.raises(HTTPException) as exc_info: + await _call(handler, "s1", max_iterations=3) + + assert exc_info.value.status_code == 429 + assert await raw.get("ns:{session_iterations:s1}:count") == b"4" + + +@pytest.mark.asyncio +async def test_the_first_iteration_sets_the_ttl(): + server = fakeredis.FakeServer() + raw = fakeredis.FakeAsyncRedis(server=server) + + await _call(_redis_handler(server), "s1", max_iterations=3) + + assert 3_590_000 < await raw.pttl("ns:{session_iterations:s1}:count") <= 3_600_000 + + +@pytest.mark.asyncio +async def test_an_unreachable_redis_counts_iterations_in_memory(): + server = fakeredis.FakeServer() + handler = _redis_handler(server) + server.connected = False + + await _call(handler, "s1", max_iterations=1) + with pytest.raises(HTTPException): + await _call(handler, "s1", max_iterations=1)