From e5d96f875082b1dd82c57d15802ac34b4b2a9cb3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 24 Feb 2026 19:47:26 -0800 Subject: [PATCH] feat(proxy): support agent-level max_iterations config max_iterations is now resolved in priority order: 1. Agent config (litellm_params.max_iterations via agent registry) 2. Key metadata (fallback) Co-Authored-By: Claude Opus 4.6 --- docs/my-website/docs/proxy/max_iterations.md | 36 +++++---- litellm/proxy/hooks/max_iterations_limiter.py | 74 +++++++++++++++++-- .../hooks/test_max_iterations_limiter.py | 50 +++++++++++++ 3 files changed, 139 insertions(+), 21 deletions(-) diff --git a/docs/my-website/docs/proxy/max_iterations.md b/docs/my-website/docs/proxy/max_iterations.md index a8444891ffe..d8dbd5617cf 100644 --- a/docs/my-website/docs/proxy/max_iterations.md +++ b/docs/my-website/docs/proxy/max_iterations.md @@ -7,18 +7,32 @@ Limit the number of LLM calls an agentic loop can make per session. Callers send ## Quick Start -### 1. Set `max_iterations` on a key +### 1. Set `max_iterations` on an agent + +Set `max_iterations` in the agent's `litellm_params` when creating the agent: + +```bash +curl -L -X POST 'http://0.0.0.0:4000/agents' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "agent_name": "my-agent", + "litellm_params": {"max_iterations": 25}, + "agent_card_params": {"name": "my-agent", "url": "http://agent:8000"} +}' +``` + +You can also set it per key as a fallback (agent config takes priority): ```bash curl -L -X POST 'http://0.0.0.0:4000/key/generate' \ -H 'Authorization: Bearer sk-1234' \ --H 'Content-Type: application/json' \ -d '{"metadata": {"max_iterations": 25}}' ``` ### 2. Send requests with `session_id` -Include the same `session_id` on every call in the agent loop via the `x-litellm-session-id` header or `metadata.session_id`. +Include the same `session_id` on every call in the agent loop via `x-litellm-session-id` header or `metadata.session_id`. @@ -56,17 +70,13 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ Works on all proxy endpoints: `/v1/chat/completions`, `/v1/responses`, `/v1/messages`, `/a2a/{agent_id}`. -## Configuration +## Priority Order -Set `max_iterations` in key metadata via `/key/generate` or `/key/update`: +`max_iterations` is resolved in this order: -```bash -# Update existing key -curl -L -X POST 'http://0.0.0.0:4000/key/update' \ --H 'Authorization: Bearer sk-1234' \ --d '{"key": "sk-existing-key", "metadata": {"max_iterations": 50}}' -``` +1. **Agent config** — `litellm_params.max_iterations` (looked up via `agent_id` in request metadata) +2. **Key metadata** — `metadata.max_iterations` (set via `/key/generate` or `/key/update`) -Session counters auto-expire after 1 hour (configurable via `LITELLM_MAX_ITERATIONS_TTL` env var in seconds). +## Settings -Works across multiple proxy instances via Redis. +Session counters auto-expire after 1 hour (configurable via `LITELLM_MAX_ITERATIONS_TTL` env var in seconds). Works across multiple proxy instances via Redis. diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 8d481f6b261..cd7fd10cbf1 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -4,7 +4,11 @@ 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. +exceeds `max_iterations`, returns 429. + +max_iterations is resolved in priority order: + 1. Agent config: agent's litellm_params.max_iterations (looked up via agent_id) + 2. Key metadata: user_api_key_dict.metadata.max_iterations Works across multiple proxy instances via DualCache (in-memory + Redis). Follows the same pattern as parallel_request_limiter_v3.py. @@ -51,9 +55,9 @@ 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} + Configuration (checked in priority order): + - Agent-level: agent's litellm_params.max_iterations (via agent registry) + - Key-level: key metadata.max_iterations (via /key/generate or /key/update) - session_id: sent by caller via x-litellm-session-id header or metadata.session_id in request body @@ -100,8 +104,8 @@ class _PROXY_MaxIterationsHandler(CustomLogger): if session_id is None: return None - # Extract max_iterations from key metadata - max_iterations = self._get_max_iterations(user_api_key_dict) + # Extract max_iterations: agent config first, then key metadata + max_iterations = self._get_max_iterations(data, user_api_key_dict) if max_iterations is None: return None @@ -149,15 +153,69 @@ class _PROXY_MaxIterationsHandler(CustomLogger): return None def _get_max_iterations( - self, user_api_key_dict: UserAPIKeyAuth + self, data: dict, user_api_key_dict: UserAPIKeyAuth ) -> Optional[int]: - """Extract max_iterations from key metadata.""" + """ + Extract max_iterations, checking agent config first, then key metadata. + + Priority: + 1. Agent litellm_params.max_iterations (via agent_id in request metadata) + 2. Key metadata.max_iterations + """ + # 1. Check agent config via agent_id + agent_id = self._get_agent_id(data) + if agent_id is not None: + agent_max = self._get_max_iterations_from_agent(agent_id) + if agent_max is not None: + return agent_max + + # 2. Fallback to 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 _get_agent_id(self, data: dict) -> Optional[str]: + """Extract agent_id from request metadata.""" + metadata = data.get("metadata") or {} + agent_id = metadata.get("agent_id") + if agent_id is not None: + return str(agent_id) + + litellm_metadata = data.get("litellm_metadata") or {} + agent_id = litellm_metadata.get("agent_id") + if agent_id is not None: + return str(agent_id) + + return None + + def _get_max_iterations_from_agent(self, agent_id: str) -> Optional[int]: + """Look up max_iterations from agent's litellm_params in the registry.""" + try: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + + agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) + if agent is None: + agent = global_agent_registry.get_agent_by_name( + agent_name=agent_id + ) + if agent is None or agent.litellm_params is None: + return None + + max_iterations = agent.litellm_params.get("max_iterations") + if max_iterations is not None: + return int(max_iterations) + except Exception as e: + verbose_proxy_logger.debug( + "MaxIterationsHandler: Could not look up agent %s: %s", + agent_id, + str(e), + ) + return None + def _make_cache_key(self, session_id: str) -> str: """ Create cache key for session iteration counter. diff --git a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py index deb1c483b87..0d7b56bf1a9 100644 --- a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py @@ -4,8 +4,11 @@ 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 +- Agent-level max_iterations takes priority over key-level """ +from unittest.mock import MagicMock, patch + import pytest from fastapi import HTTPException @@ -104,3 +107,50 @@ async def test_max_iterations_different_sessions_independent(): data={"metadata": {"session_id": "session-B"}}, call_type="", ) + + +@pytest.mark.asyncio +async def test_max_iterations_agent_level_config(): + """ + Test that max_iterations from agent config takes priority over key metadata. + + - Agent has max_iterations=2 in litellm_params + - Key has max_iterations=100 in metadata + - Agent limit (2) should be enforced, not key limit (100) + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + # Key allows 100 iterations + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-agent", metadata={"max_iterations": 100} + ) + + # Mock agent registry to return an agent with max_iterations=2 + mock_agent = MagicMock() + mock_agent.litellm_params = {"max_iterations": 2} + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + # 2 calls succeed (agent limit) + 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-agent", "agent_id": "agent-123"}}, + call_type="", + ) + + # 3rd call fails (agent limit of 2 exceeded, not key limit of 100) + 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-agent", "agent_id": "agent-123"}}, + call_type="", + ) + assert exc_info.value.status_code == 429