From 84874012ad769c6edc70072afabc7ec641677f2d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 24 Feb 2026 19:37:21 -0800 Subject: [PATCH] feat(proxy): add max_iterations limiter for agent session loops Adds a new proxy hook that enforces a per-session cap on the number of LLM calls an agentic loop can make. Callers send a session_id with each request, and the hook counts calls per session, returning 429 when the configured max_iterations limit is exceeded. - Uses Redis Lua script for atomic increment (multi-instance safe) - Falls back to in-memory cache when Redis unavailable - Follows parallel_request_limiter_v3 pattern - Configurable via key metadata: {"max_iterations": 25} - Session counters auto-expire via TTL (default 1hr) Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/hooks/max_iterations_limiter.py | 208 ++++++++++++++++++ .../hooks/test_max_iterations_limiter.py | 106 +++++++++ 2 files changed, 314 insertions(+) create mode 100644 litellm/proxy/hooks/max_iterations_limiter.py create mode 100644 tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py new file mode 100644 index 00000000000..8d481f6b261 --- /dev/null +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -0,0 +1,208 @@ +""" +Max Iterations Limiter for LiteLLM Proxy. + +Enforces a per-session cap on the number of LLM calls an agentic loop can make. +Callers send a `session_id` with each request (via `x-litellm-session-id` header +or `metadata.session_id`), and this hook counts calls per session. When the count +exceeds `max_iterations` (configured in key/team metadata), returns 429. + +Works across multiple proxy instances via DualCache (in-memory + Redis). +Follows the same pattern as parallel_request_limiter_v3.py. +""" + +import os +from typing import TYPE_CHECKING, Any, Optional, Union + +from fastapi import HTTPException + +from litellm import DualCache +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + +if TYPE_CHECKING: + from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + + InternalUsageCache = _InternalUsageCache +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 = """ +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 = 3600 + + +class _PROXY_MaxIterationsHandler(CustomLogger): + """ + Pre-call hook that enforces max_iterations per session. + + Configuration: + - max_iterations: set in key metadata via /key/generate or /key/update + e.g. metadata={"max_iterations": 25} + - session_id: sent by caller via x-litellm-session-id header or + metadata.session_id in request body + + Cache key pattern: + {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. + """ + + 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 + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ) -> Optional[Union[Exception, str, dict]]: + """ + Check session iteration count before making the API call. + + Extracts session_id from request metadata and max_iterations from + key metadata. If the session has exceeded max_iterations, raises 429. + """ + # Extract session_id from request data + session_id = self._get_session_id(data) + if session_id is None: + return None + + # Extract max_iterations from key metadata + max_iterations = self._get_max_iterations(user_api_key_dict) + if max_iterations is None: + return None + + verbose_proxy_logger.debug( + "MaxIterationsHandler: session_id=%s, max_iterations=%s", + session_id, + max_iterations, + ) + + # Increment and check + cache_key = self._make_cache_key(session_id) + current_count = await self._increment_and_get(cache_key) + + if current_count > max_iterations: + raise HTTPException( + status_code=429, + detail=( + f"Max iterations exceeded for session {session_id}. " + f"Current count: {current_count}, max_iterations: {max_iterations}." + ), + ) + + verbose_proxy_logger.debug( + "MaxIterationsHandler: session_id=%s, count=%s/%s", + session_id, + current_count, + max_iterations, + ) + + return None + + def _get_session_id(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata.""" + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + + # Also check litellm_metadata (used for /thread and /assistant endpoints) + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id is not None: + return str(session_id) + + return None + + def _get_max_iterations( + self, user_api_key_dict: UserAPIKeyAuth + ) -> Optional[int]: + """Extract max_iterations from key metadata.""" + metadata = user_api_key_dict.metadata or {} + max_iterations = metadata.get("max_iterations") + if max_iterations is not None: + return int(max_iterations) + return None + + def _make_cache_key(self, session_id: str) -> str: + """ + Create cache key for session iteration counter. + + 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, 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 = 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) + + async def _in_memory_increment(self, cache_key: str) -> int: + """Increment counter in in-memory cache with TTL.""" + current = await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) + new_value = (int(current) if current is not None else 0) + 1 + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=new_value, + ttl=self.ttl, + litellm_parent_otel_span=None, + local_only=True, + ) + return new_value diff --git a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py new file mode 100644 index 00000000000..deb1c483b87 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py @@ -0,0 +1,106 @@ +""" +Unit Tests for the max iterations limiter for the proxy. + +Tests that session-scoped iteration counting works correctly: +- Enforces max_iterations per session_id +- Different sessions have independent counters +""" + +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler +from litellm.proxy.utils import InternalUsageCache + + +@pytest.mark.asyncio +async def test_max_iterations_basic_enforcement(): + """ + Test that max_iterations is enforced per session_id. + + - 3 requests with the same session_id should succeed when max_iterations=3 + - 4th request should raise 429 + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-1234", metadata={"max_iterations": 3} + ) + + # First 3 requests should succeed + for i in range(3): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + + # 4th request should fail with 429 + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "max_iterations" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_max_iterations_different_sessions_independent(): + """ + Test that different session_ids have independent iteration counters. + + - Session A and Session B each get their own max_iterations budget + - Exhausting Session A does not affect Session B + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-5678", metadata={"max_iterations": 2} + ) + + # Session A: 2 calls succeed + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + + # Session B: 2 calls succeed (independent counter) + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + + # Session A: 3rd call fails + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + # Session B: 3rd call also fails + with pytest.raises(HTTPException): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + )