From ba3ecf9d9e1d461fe6cca355b8f3df68d6aa62f3 Mon Sep 17 00:00:00 2001 From: AlisinaDevelo Date: Sun, 27 Sep 2026 08:31:18 +0200 Subject: [PATCH] fix: preserve per-agent session limits on upgrade --- .../hooks/max_budget_per_session_limiter.py | 236 +++++++++++----- litellm/proxy/hooks/max_iterations_limiter.py | 211 +++++++++++---- .../test_max_budget_per_session_limiter.py | 251 +++++++++++++++--- .../hooks/test_max_iterations_limiter.py | 154 ++++++++++- .../test_proxy_rate_limit_provider_field.py | 4 +- .../unit/test_rate_limit_error_unification.py | 1 + 6 files changed, 704 insertions(+), 153 deletions(-) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 0ac1547693e..6a3e66ff842 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -14,10 +14,12 @@ Works across multiple proxy instances via DualCache (in-memory + Redis). Follows the same pattern as max_iterations_limiter.py. """ +import asyncio import json import logging import os -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any, Final, cast from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -36,21 +38,71 @@ else: InternalUsageCache = Any -# Redis Lua script for atomic float increment with TTL. -# INCRBYFLOAT returns the new value as a string. -# Only sets EXPIRE on first call (when prior value was nil). +# Redis Lua scripts keep the legacy aggregate and new per-agent counters in +# sync. All keys use the legacy session hash tag so this also works on Redis +# Cluster. Old proxy instances continue to update the aggregate key. MAX_BUDGET_SESSION_INCREMENT_SCRIPT: Final = """ -local key = KEYS[1] -local amount = ARGV[1] +local legacy_key = KEYS[1] +local total_new_key = KEYS[2] +local agent_key = KEYS[3] +local amount = tonumber(ARGV[1]) local ttl = tonumber(ARGV[2]) -local existed = redis.call('EXISTS', key) -local new_val = redis.call('INCRBYFLOAT', key, amount) -if existed == 0 then - redis.call('EXPIRE', key, ttl) +if redis.call('EXISTS', total_new_key) == 0 then + if redis.call('EXISTS', agent_key) == 1 then + return redis.error_reply('agent session spend exists without migration total') + end + + if redis.call('EXISTS', legacy_key) == 0 then + redis.call('SET', legacy_key, '0') + redis.call('PEXPIRE', legacy_key, ttl * 1000) + end + + local legacy_ttl = redis.call('PTTL', legacy_key) + if legacy_ttl == -2 then + return redis.error_reply('legacy session spend expired during migration') + end + + redis.call('SET', total_new_key, '0') + if legacy_ttl >= 0 then + redis.call('PEXPIRE', total_new_key, legacy_ttl + 1000) + end end -return new_val +if redis.call('EXISTS', legacy_key) == 0 then + return redis.error_reply('legacy session spend expired before agent scope') +end + +local legacy_value = tonumber(redis.call('INCRBYFLOAT', legacy_key, amount)) +local total_new_value = tonumber(redis.call('INCRBYFLOAT', total_new_key, amount)) +local agent_existed = redis.call('EXISTS', agent_key) +local agent_value = tonumber(redis.call('INCRBYFLOAT', agent_key, amount)) +if agent_existed == 0 then + local migration_ttl = redis.call('PTTL', total_new_key) + if migration_ttl >= 0 then + redis.call('PEXPIRE', agent_key, migration_ttl + 1000) + end +end + +return tostring(math.max(legacy_value - total_new_value, 0) + agent_value) +""" + +MAX_BUDGET_SESSION_GET_AGENT_SPEND_SCRIPT: Final = """ +local legacy_value = tonumber(redis.call('GET', KEYS[1])) or 0 +local total_new_value = tonumber(redis.call('GET', KEYS[2])) +local agent_value = tonumber(redis.call('GET', KEYS[3])) or 0 + +if total_new_value == nil then + if redis.call('EXISTS', KEYS[3]) == 1 then + return redis.error_reply('agent session spend exists without migration total') + end + return tostring(math.max(legacy_value, agent_value)) +end +if redis.call('EXISTS', KEYS[1]) == 0 then + return redis.error_reply('legacy session spend expired before agent scope') +end + +return tostring(math.max(legacy_value - total_new_value, 0) + agent_value) """ # Default TTL for session budget counters (1 hour) @@ -65,11 +117,14 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): - max_budget_per_session: dollar cap per agent and session_id Cache key pattern: - {agent_session_budget:[,]}:spend + {session_budget:}:agent::spend """ def __init__(self, internal_usage_cache: InternalUsageCache): self.internal_usage_cache = internal_usage_cache + self._local_lock = asyncio.Lock() + self.increment_script: Callable[..., Awaitable[object]] | None = None + self.get_agent_spend_script: Callable[..., Awaitable[object]] | None = None self.ttl = int( os.getenv( "LITELLM_MAX_BUDGET_PER_SESSION_TTL", @@ -78,11 +133,18 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) 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_BUDGET_SESSION_INCREMENT_SCRIPT + self.increment_script = cast( + Callable[..., Awaitable[object]], + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_BUDGET_SESSION_INCREMENT_SCRIPT + ), + ) + self.get_agent_spend_script = cast( + Callable[..., Awaitable[object]], + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_BUDGET_SESSION_GET_AGENT_SPEND_SCRIPT + ), ) - else: - self.increment_script = None async def async_pre_call_hook( self, @@ -104,8 +166,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return None max_budget = float(max_budget) - cache_key: Final = self._make_cache_key(session_id, agent_id) - current_spend: Final = await self._get_current_spend(cache_key) + current_spend: Final = await self._get_agent_spend(session_id, agent_id) verbose_proxy_logger.debug( "MaxBudgetPerSessionHandler: session_id=%s, spend=%.4f, max=%.2f", @@ -131,7 +192,9 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ - After a successful LLM call, increment the session spend by the response cost. + Record every successful agent call so limits added later still see the + session's accumulated spend. The pre-call hook enforces a cap only when + one is configured. """ try: litellm_params: Final = kwargs.get("litellm_params") or {} @@ -152,17 +215,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): if agent is None: return - agent_litellm_params: Final = agent.litellm_params or {} - max_budget: Final = agent_litellm_params.get("max_budget_per_session") - if max_budget is None: - return - response_cost: Final = kwargs.get("response_cost") or 0.0 if response_cost <= 0: return - cache_key: Final = self._make_cache_key(str(session_id), agent.agent_id) - await self._increment_spend(cache_key, float(response_cost)) + await self._increment_agent_spend(str(session_id), agent.agent_id, float(response_cost)) verbose_proxy_logger.debug( "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", @@ -174,6 +231,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): "MaxBudgetPerSessionHandler: error in async_log_success_event: %s", str(e), ) + raise def _get_session_id(self, data: dict) -> str | None: """Extract session_id from request metadata.""" @@ -210,65 +268,115 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): def _make_cache_key(self, session_id: str, agent_id: str) -> str: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - scope: Final = json.dumps((global_agent_registry.stable_agent_id(agent_id), session_id), separators=(",", ":")) - return f"{{agent_session_budget:{scope}}}:spend" + stable_agent_id: Final = json.dumps(global_agent_registry.stable_agent_id(agent_id), separators=(",", ":")) + return f"{{session_budget:{session_id}}}:agent:{stable_agent_id}:spend" - async def _get_current_spend(self, cache_key: str) -> float: - """Read current accumulated spend for a session.""" - if self.internal_usage_cache.dual_cache.redis_cache is not None: + def _make_legacy_cache_key(self, session_id: str) -> str: + return f"{{session_budget:{session_id}}}:spend" + + def _make_total_new_cache_key(self, session_id: str) -> str: + return f"{{session_budget:{session_id}}}:agent-scope-total" + + async def _get_agent_spend(self, session_id: str, agent_id: str) -> float: + legacy_key: Final = self._make_legacy_cache_key(session_id) + total_new_key: Final = self._make_total_new_cache_key(session_id) + agent_key: Final = self._make_cache_key(session_id, agent_id) + if self.get_agent_spend_script is not None: try: - result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key) - if result is not None: + result: Final[object] = await self.get_agent_spend_script( + keys=[legacy_key, total_new_key, agent_key], + args=[], + ) + if isinstance(result, (int, float, str, bytes)): return float(result) - return 0.0 + raise TypeError(f"Unexpected Redis spend result: {type(result).__name__}") except Exception as e: log_redis_failure( verbose_proxy_logger, logging.WARNING, - "MaxBudgetPerSessionHandler: Redis GET failed, falling back to in-memory", + "MaxBudgetPerSessionHandler: Redis agent spend read failed", e, ) + raise - result = await self.internal_usage_cache.async_get_cache( - key=cache_key, - litellm_parent_otel_span=None, - local_only=True, + legacy_value: Final = await self._get_local_spend(legacy_key) + total_new_value: Final = await self._get_local_spend(total_new_key) + agent_value: Final = await self._get_local_spend(agent_key) + if legacy_value is None: + if total_new_value is not None: + raise RuntimeError("Legacy session spend expired before agent scope") + return float(agent_value or 0.0) + if total_new_value is None: + return max(float(legacy_value), float(agent_value or 0.0)) + return max(float(legacy_value) - float(total_new_value), 0.0) + float(agent_value or 0.0) + + async def _get_local_spend(self, cache_key: str) -> float | None: + result: Final[object | None] = cast( + object | None, + await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ), ) - if result is not None: + if isinstance(result, (int, float, str, bytes)): return float(result) - return 0.0 + return None - async def _increment_spend(self, cache_key: str, amount: float) -> float: - """Atomically increment the session spend and return the new value.""" + async def _increment_agent_spend(self, session_id: str, agent_id: str, amount: float) -> float: + legacy_key: Final = self._make_legacy_cache_key(session_id) + total_new_key: Final = self._make_total_new_cache_key(session_id) + agent_key: Final = self._make_cache_key(session_id, agent_id) if self.increment_script is not None: try: - result: Final = await self.increment_script( - keys=[cache_key], + result: Final[object] = await self.increment_script( + keys=[legacy_key, total_new_key, agent_key], args=[str(amount), self.ttl], ) - return float(result) + if isinstance(result, (int, float, str, bytes)): + return float(result) + raise TypeError(f"Unexpected Redis spend result: {type(result).__name__}") except Exception as e: log_redis_failure( verbose_proxy_logger, logging.WARNING, - "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory", + "MaxBudgetPerSessionHandler: Redis migration increment failed; refusing an unsafe retry", e, ) + raise - return await self._in_memory_increment_spend(cache_key, amount) - - async def _in_memory_increment_spend(self, cache_key: str, amount: float) -> float: - current: Final = await self.internal_usage_cache.async_get_cache( - key=cache_key, - litellm_parent_otel_span=None, - local_only=True, - ) - new_value: Final = (float(current) if current is not None else 0.0) + amount - 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 + async with self._local_lock: + legacy_value = await self._get_local_spend(legacy_key) + total_new_value = await self._get_local_spend(total_new_key) + agent_value = await self._get_local_spend(agent_key) + if legacy_value is None: + if total_new_value is not None: + raise RuntimeError("Legacy session spend expired before agent scope") + legacy_value = 0.0 + total_new_value = total_new_value or 0.0 + agent_value = agent_value or 0.0 + new_legacy: Final = float(legacy_value) + amount + new_total_new: Final = float(total_new_value) + amount + new_agent: Final = float(agent_value) + amount + await self.internal_usage_cache.async_set_cache( + key=legacy_key, + value=new_legacy, + ttl=self.ttl if legacy_value == 0.0 else None, + litellm_parent_otel_span=None, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=total_new_key, + value=new_total_new, + ttl=self.ttl + 1 if total_new_value == 0.0 else None, + litellm_parent_otel_span=None, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=agent_key, + value=new_agent, + ttl=self.ttl + 1 if agent_value == 0.0 else None, + litellm_parent_otel_span=None, + local_only=True, + ) + return max(new_legacy - new_total_new, 0.0) + new_agent diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index eb27f65fcc4..9288a71c611 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -10,9 +10,11 @@ Works across multiple proxy instances via DualCache (in-memory + Redis). Follows the same pattern as parallel_request_limiter_v3.py. """ +import asyncio import json import os -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any, Final, cast from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -30,19 +32,59 @@ 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). +# Redis Lua scripts keep the legacy aggregate and new per-agent counters in +# sync. All keys use the legacy session hash tag so this also works on Redis +# Cluster. Old proxy instances continue to update the aggregate key. 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) +local legacy_key = KEYS[1] +if #KEYS == 1 then + local current = redis.call('INCR', legacy_key) + if current == 1 then + redis.call('EXPIRE', legacy_key, tonumber(ARGV[1])) + end + return current end -return current +local total_new_key = KEYS[2] +local agent_key = KEYS[3] +local ttl = tonumber(ARGV[1]) +if redis.call('EXISTS', total_new_key) == 0 then + if redis.call('EXISTS', agent_key) == 1 then + return redis.error_reply('agent session count exists without migration total') + end + + if redis.call('EXISTS', legacy_key) == 0 then + redis.call('SET', legacy_key, '0') + redis.call('PEXPIRE', legacy_key, ttl * 1000) + end + + local legacy_ttl = redis.call('PTTL', legacy_key) + if legacy_ttl == -2 then + return redis.error_reply('legacy session count expired during migration') + end + + redis.call('SET', total_new_key, '0') + if legacy_ttl >= 0 then + redis.call('PEXPIRE', total_new_key, legacy_ttl + 1000) + end +end + +if redis.call('EXISTS', legacy_key) == 0 then + return redis.error_reply('legacy session count expired before agent scope') +end + +local legacy_value = redis.call('INCR', legacy_key) +local total_new_value = redis.call('INCR', total_new_key) +local agent_existed = redis.call('EXISTS', agent_key) +local agent_value = redis.call('INCR', agent_key) +if agent_existed == 0 then + local migration_ttl = redis.call('PTTL', total_new_key) + if migration_ttl >= 0 then + redis.call('PEXPIRE', agent_key, migration_ttl + 1000) + end +end + +return math.max(legacy_value - total_new_value, 0) + agent_value """ # Default TTL for session iteration counters (1 hour) @@ -61,26 +103,28 @@ class _PROXY_MaxIterationsHandler(CustomLogger): metadata.session_id in request body Cache key pattern: - {agent_session_iterations:[,]}:count + {session_iterations:}:agent::count Without an agent, retains {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. + Uses Redis Lua scripts for atomic increments when Redis is configured. + Uses process-local memory only when Redis is not configured; Redis + errors propagate so a failed shared counter cannot silently bypass the + limit. """ def __init__(self, internal_usage_cache: InternalUsageCache): self.internal_usage_cache = internal_usage_cache + self._local_lock = asyncio.Lock() + self.increment_script: Callable[..., Awaitable[object]] | None = None 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 + self.increment_script = cast( + Callable[..., Awaitable[object]], + 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, @@ -111,8 +155,10 @@ class _PROXY_MaxIterationsHandler(CustomLogger): ) # Increment and check - cache_key: Final = self._make_cache_key(session_id, user_api_key_dict.agent_id) - current_count: Final = await self._increment_and_get(cache_key) + if user_api_key_dict.agent_id is None: + current_count = await self._increment_legacy_and_get(self._make_legacy_cache_key(session_id)) + else: + current_count = await self._increment_agent_and_get(session_id, user_api_key_dict.agent_id) if current_count > max_iterations: resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) @@ -177,52 +223,109 @@ class _PROXY_MaxIterationsHandler(CustomLogger): """ Create cache key for session iteration counter. - The Redis hash tag includes both identities when an agent is configured. - Keys without an agent retain the legacy session scope. + Agent-scoped counters share the legacy session hash tag so migration + scripts can atomically update both scopes on Redis Cluster. """ if agent_id is None: return f"{{session_iterations:{session_id}}}:count" from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - scope: Final = json.dumps((global_agent_registry.stable_agent_id(agent_id), session_id), separators=(",", ":")) - return f"{{agent_session_iterations:{scope}}}:count" + stable_agent_id: Final = json.dumps(global_agent_registry.stable_agent_id(agent_id), separators=(",", ":")) + return f"{{session_iterations:{session_id}}}:agent:{stable_agent_id}:count" - async def _increment_and_get(self, cache_key: str) -> int: - """ - Atomically increment the session counter and return the new value. + def _make_legacy_cache_key(self, session_id: str) -> str: + return f"{{session_iterations:{session_id}}}:count" - Tries Redis first (via registered Lua script for atomicity across - instances), falls back to in-memory cache. - """ + def _make_total_new_cache_key(self, session_id: str) -> str: + return f"{{session_iterations:{session_id}}}:agent-scope-total" + + async def _get_local_count(self, cache_key: str) -> int | None: + local_result: Final[object | None] = cast( + object | None, + await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ), + ) + if isinstance(local_result, (int, float, str, bytes)): + return int(local_result) + return None + + async def _increment_legacy_and_get(self, cache_key: str) -> int: + if self.increment_script is not None: + result: Final[object] = await self.increment_script( + keys=[cache_key], + args=[self.ttl], + ) + if isinstance(result, (int, float, str, bytes)): + return int(result) + raise TypeError(f"Unexpected Redis iteration result: {type(result).__name__}") + + async with self._local_lock: + current: Final = await self._get_local_count(cache_key) + new_value: Final = (current or 0) + 1 + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=new_value, + ttl=self.ttl if current is None else None, + litellm_parent_otel_span=None, + local_only=True, + ) + return new_value + + async def _increment_agent_and_get(self, session_id: str, agent_id: str) -> int: + legacy_key: Final = self._make_legacy_cache_key(session_id) + total_new_key: Final = self._make_total_new_cache_key(session_id) + agent_key: Final = self._make_cache_key(session_id, agent_id) if self.increment_script is not None: try: - result: Final = await self.increment_script( - keys=[cache_key], + result: Final[object] = await self.increment_script( + keys=[legacy_key, total_new_key, agent_key], args=[self.ttl], ) - return int(result) + if isinstance(result, (int, float, str, bytes)): + return int(result) + raise TypeError(f"Unexpected Redis iteration result: {type(result).__name__}") except Exception as e: verbose_proxy_logger.warning( - "MaxIterationsHandler: Redis failed, falling back to in-memory: %s", + "MaxIterationsHandler: Redis migration increment failed; refusing an unsafe retry: %s", str(e), ) + raise - # 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: Final = await self.internal_usage_cache.async_get_cache( - key=cache_key, - litellm_parent_otel_span=None, - local_only=True, - ) - new_value: Final = (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 + async with self._local_lock: + legacy_value = await self._get_local_count(legacy_key) + total_new_value = await self._get_local_count(total_new_key) + agent_value = await self._get_local_count(agent_key) + if legacy_value is None: + if total_new_value is not None: + raise RuntimeError("Legacy session count expired before agent scope") + legacy_value = 0 + total_new_value = total_new_value or 0 + agent_value = agent_value or 0 + new_legacy: Final = legacy_value + 1 + new_total_new: Final = total_new_value + 1 + new_agent: Final = agent_value + 1 + await self.internal_usage_cache.async_set_cache( + key=legacy_key, + value=new_legacy, + ttl=self.ttl if legacy_value == 0 else None, + litellm_parent_otel_span=None, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=total_new_key, + value=new_total_new, + ttl=self.ttl + 1 if total_new_value == 0 else None, + litellm_parent_otel_span=None, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=agent_key, + value=new_agent, + ttl=self.ttl + 1 if agent_value == 0 else None, + litellm_parent_otel_span=None, + local_only=True, + ) + return max(new_legacy - new_total_new, 0) + new_agent diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py index de4f581d4e7..5c5b84c6ed8 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py @@ -8,15 +8,17 @@ Tests that session-scoped budget tracking works correctly: - Requests without agent_id pass through """ -import logging +import asyncio +import os +import socket +import uuid from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from fastapi import HTTPException from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.hooks.max_budget_per_session_limiter import ( @@ -35,6 +37,13 @@ def _make_mock_agent(max_budget_per_session: float, agent_id: str = "agent-budge ) +def _redis_port_for_migration_test() -> int | None: + port = int(os.getenv("LITELLM_TEST_REDIS_PORT", "6379")) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + return port if sock.connect_ex(("127.0.0.1", port)) == 0 else None + + @pytest.mark.asyncio async def test_budget_per_session_under_budget_passes(): """ @@ -79,8 +88,7 @@ async def test_budget_per_session_exceeds_budget(): ) session_id = "session-over-budget" - cache_key = handler._make_cache_key(session_id, "agent-budget-123") - await handler._increment_spend(cache_key, 1.50) + await handler._increment_agent_spend(session_id, "agent-budget-123", 1.50) mock_agent = _make_mock_agent(max_budget_per_session=1.0) @@ -113,8 +121,7 @@ async def test_budget_per_session_independent_sessions(): agent_id="agent-budget-123", ) - cache_key_a = handler._make_cache_key("session-A", "agent-budget-123") - await handler._increment_spend(cache_key_a, 3.0) + await handler._increment_agent_spend("session-A", "agent-budget-123", 3.0) mock_agent = _make_mock_agent(max_budget_per_session=2.0) @@ -245,35 +252,215 @@ async def test_legacy_agent_id_shares_the_registered_agents_budget() -> None: assert "Current spend: $1.0000" in str(rejected.value.detail) -class _OpenBreakerRedis: - def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker +@pytest.mark.asyncio +async def test_agent_budget_keeps_spend_from_an_existing_session() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(cache)) + session_id: Final = "existing-session" + legacy_key: Final = handler._make_legacy_cache_key(session_id) + await cache.async_set_cache(key=legacy_key, value=1.0) + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(1.0)) - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + with pytest.raises(HTTPException) as rejected: + await handler.async_pre_call_hook( + UserAPIKeyAuth(agent_id="agent-budget-123"), + cache, + {"metadata": {"session_id": session_id}}, + "", + ) - @_redis_circuit_breaker_guard - async def async_get_cache(self, key, **kwargs): - raise AssertionError("never reached") - - def async_register_script(self, script): - @_redis_circuit_breaker_guard - async def refused(_self, keys, args): - raise AssertionError("never reached") - - return lambda keys, args: refused(self, keys, args) + assert rejected.value.status_code == 429 + assert "Current spend: $1.0000" in str(rejected.value.detail) @pytest.mark.asyncio -async def test_an_open_circuit_breaker_reads_session_spend_locally_without_a_warning(caplog): - cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double - handler = _PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(cache)) - caplog.clear() +async def test_agent_budget_carries_existing_spend_into_its_counter() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(cache)) + session_id: Final = "existing-session" + await cache.async_set_cache(key=handler._make_legacy_cache_key(session_id), value=0.75) + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(1.0)) + kwargs: Final = { + "litellm_params": {"metadata": {"session_id": session_id, "agent_id": "agent-budget-123"}}, + "response_cost": 0.1, + } - with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): - spend = await handler._get_current_spend("{session_budget:quiet}:spend") + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + await handler.async_log_success_event(kwargs, None, None, None) - assert spend == 0.0 - assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] - assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) + spend: Final = await handler._get_agent_spend(session_id, "agent-budget-123") + assert spend == pytest.approx(0.85) + + +@pytest.mark.asyncio +async def test_agent_budget_tracks_old_and_new_pods_without_mixing_new_agent_spend() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(cache)) + session_id: Final = "mixed-rollout-session" + legacy_key: Final = handler._make_legacy_cache_key(session_id) + await cache.async_set_cache(key=legacy_key, value=0.5) + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(3.0, "agent-a")) + registry.register_agent(_make_mock_agent(3.0, "agent-b")) + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + await handler._increment_agent_spend(session_id, "agent-a", 0.4) + assert await handler._get_agent_spend(session_id, "agent-a") == pytest.approx(0.9) + assert await handler._get_agent_spend(session_id, "agent-b") == pytest.approx(0.5) + + # An old pod still writes only the legacy aggregate. + await cache.async_set_cache(key=legacy_key, value=1.1) + await handler._increment_agent_spend(session_id, "agent-a", 0.1) + await handler._increment_agent_spend(session_id, "agent-b", 0.1) + + assert await handler._get_agent_spend(session_id, "agent-a") == pytest.approx(1.2) + assert await handler._get_agent_spend(session_id, "agent-b") == pytest.approx(0.8) + + +@pytest.mark.asyncio +async def test_agent_spend_is_recorded_before_a_budget_is_configured() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(cache)) + session_id: Final = "unlimited-session" + agent: Final = _make_mock_agent(1.0) + agent.litellm_params = {} + registry: Final = AgentRegistry() + registry.register_agent(agent) + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + await handler.async_log_success_event( + { + "litellm_params": { + "metadata": {"session_id": session_id, "agent_id": agent.agent_id}, + }, + "response_cost": 1.25, + }, + None, + None, + None, + ) + agent.litellm_params["max_budget_per_session"] = 1.0 + + with pytest.raises(HTTPException) as rejected: + await handler.async_pre_call_hook( + UserAPIKeyAuth(agent_id=agent.agent_id), + cache, + {"metadata": {"session_id": session_id}}, + "", + ) + + assert rejected.value.status_code == 429 + assert "Current spend: $1.2500" in str(rejected.value.detail) + + +@pytest.mark.asyncio +async def test_agent_budget_redis_errors_are_not_retried_or_read_locally() -> None: + handler: Final = _PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(DualCache())) + increment_calls = 0 + read_calls = 0 + + async def fail_increment(**_kwargs: object) -> object: + nonlocal increment_calls + increment_calls += 1 + raise TimeoutError("reply timed out after Redis may have applied the script") + + async def fail_read(**_kwargs: object) -> object: + nonlocal read_calls + read_calls += 1 + raise TimeoutError("Redis read timed out") + + handler.increment_script = fail_increment + handler.get_agent_spend_script = fail_read + with patch.object(handler, "_get_local_spend", new=AsyncMock(side_effect=AssertionError("must fail closed"))): + with pytest.raises(TimeoutError): + await handler._increment_agent_spend("uncertain-session", "agent-a", 0.1) + with pytest.raises(TimeoutError): + await handler._get_agent_spend("uncertain-session", "agent-a") + + assert increment_calls == 1 + assert read_calls == 1 + + +@pytest.mark.asyncio +@pytest.mark.skipif(_redis_port_for_migration_test() is None, reason="requires local Redis for Lua migration path") +async def test_redis_budget_migration_is_atomic_and_preserves_session_ttl() -> None: + from litellm.caching.redis_cache import RedisCache + + port: Final = _redis_port_for_migration_test() + assert port is not None + redis: Final = RedisCache(host="127.0.0.1", port=port) + redis_client: Final = redis.init_async_client() + handler: Final = _PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(DualCache(redis_cache=redis))) + handler.ttl = 30 + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(10.0, "agent-a")) + registry.register_agent(_make_mock_agent(10.0, "agent-b")) + session_id: Final = f"budget-migration-{uuid.uuid4().hex}" + concurrent_session: Final = f"budget-concurrent-{uuid.uuid4().hex}" + keys: Final = [] + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + try: + legacy_key: Final = handler._make_legacy_cache_key(session_id) + total_new_key: Final = handler._make_total_new_cache_key(session_id) + agent_a_key: Final = handler._make_cache_key(session_id, "agent-a") + agent_b_key: Final = handler._make_cache_key(session_id, "agent-b") + keys.extend((legacy_key, total_new_key, agent_a_key, agent_b_key)) + await redis.async_set_cache(key=legacy_key, value=0.5, ttl=30) + legacy_redis_key: Final = redis.check_and_fix_namespace(legacy_key) + initial_ttl: Final = await redis_client.pttl(legacy_redis_key) + + await handler._increment_agent_spend(session_id, "agent-a", 0.4) + await redis.async_increment(legacy_key, 0.2, ttl=30) # Old pod updates only the aggregate. + await handler._increment_agent_spend(session_id, "agent-a", 0.1) + await handler._increment_agent_spend(session_id, "agent-b", 0.1) + assert await handler._get_agent_spend(session_id, "agent-a") == pytest.approx(1.2) + assert await handler._get_agent_spend(session_id, "agent-b") == pytest.approx(0.8) + + hash_tags: Final = {key[key.index("{") : key.index("}") + 1] for key in keys[:4]} + assert len(hash_tags) == 1 + sidecar_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(total_new_key)) + agent_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(agent_a_key)) + current_ttl: Final = await redis_client.pttl(legacy_redis_key) + assert current_ttl <= sidecar_ttl <= current_ttl + 1100 + assert sidecar_ttl <= agent_ttl <= sidecar_ttl + 1100 + assert current_ttl <= initial_ttl + await asyncio.sleep(0.25) + before_increment_ttl: Final = await redis_client.pttl(legacy_redis_key) + before_increment_sidecar_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(total_new_key)) + await handler._increment_agent_spend(session_id, "agent-a", 0.01) + after_increment_ttl: Final = await redis_client.pttl(legacy_redis_key) + after_increment_sidecar_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(total_new_key)) + assert after_increment_ttl <= before_increment_ttl + 50 + assert after_increment_sidecar_ttl <= before_increment_sidecar_ttl + 50 + + concurrent_legacy_key: Final = handler._make_legacy_cache_key(concurrent_session) + keys.extend( + ( + concurrent_legacy_key, + handler._make_total_new_cache_key(concurrent_session), + handler._make_cache_key(concurrent_session, "agent-a"), + handler._make_cache_key(concurrent_session, "agent-b"), + ) + ) + await redis.async_set_cache(key=concurrent_legacy_key, value=5.0, ttl=30) + await asyncio.gather( + *(handler._increment_agent_spend(concurrent_session, "agent-a", 0.01) for _ in range(100)), + *(handler._increment_agent_spend(concurrent_session, "agent-b", 0.01) for _ in range(50)), + ) + assert await handler._get_agent_spend(concurrent_session, "agent-a") == pytest.approx(6.0) + assert await handler._get_agent_spend(concurrent_session, "agent-b") == pytest.approx(5.5) + + await redis_client.pexpire(legacy_redis_key, 100) + await asyncio.sleep(0.15) + with pytest.raises(Exception, match="legacy session spend expired before agent scope"): + await handler._get_agent_spend(session_id, "agent-a") + + await redis_client.delete(redis.check_and_fix_namespace(total_new_key)) + with pytest.raises(Exception, match="agent session spend exists without migration total"): + await handler._get_agent_spend(session_id, "agent-a") + finally: + await redis_client.delete(*(redis.check_and_fix_namespace(key) for key in keys)) 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 9d512859f1b..7839b3f0cf8 100644 --- a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py @@ -6,8 +6,12 @@ Tests that session-scoped iteration counting works correctly: - Different sessions have independent counters """ +import asyncio +import os +import socket +import uuid from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from fastapi import HTTPException @@ -29,6 +33,13 @@ def _make_mock_agent(max_iterations: int, agent_id: str = "agent-test-123") -> A ) +def _redis_port_for_migration_test() -> int | None: + port = int(os.getenv("LITELLM_TEST_REDIS_PORT", "6379")) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + return port if sock.connect_ex(("127.0.0.1", port)) == 0 else None + + @pytest.mark.asyncio async def test_max_iterations_basic_enforcement(): """ @@ -229,3 +240,144 @@ async def test_legacy_agent_id_shares_the_registered_agents_iteration_limit() -> await handler.async_pre_call_hook(UserAPIKeyAuth(agent_id=legacy_id), cache, data, "") assert rejected.value.status_code == 429 assert "Current count: 3" in str(rejected.value.detail) + + +@pytest.mark.asyncio +async def test_agent_iteration_limit_keeps_count_from_an_existing_session() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxIterationsHandler(InternalUsageCache(cache)) + session_id: Final = "existing-session" + legacy_key: Final = handler._make_legacy_cache_key(session_id) + await cache.async_set_cache(key=legacy_key, value=2) + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(2)) + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + with pytest.raises(HTTPException) as rejected: + await handler.async_pre_call_hook( + UserAPIKeyAuth(agent_id="agent-test-123"), + cache, + {"metadata": {"session_id": session_id}}, + "", + ) + + assert rejected.value.status_code == 429 + assert "Current count: 3" in str(rejected.value.detail) + legacy_count: Final = await cache.async_get_cache(key=handler._make_legacy_cache_key(session_id)) + total_new_count: Final = await cache.async_get_cache(key=handler._make_total_new_cache_key(session_id)) + agent_count: Final = await cache.async_get_cache(key=handler._make_cache_key(session_id, "agent-test-123")) + assert legacy_count - total_new_count + agent_count == 3 + + +@pytest.mark.asyncio +async def test_agent_iteration_counter_tracks_old_and_new_pods_independently() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxIterationsHandler(InternalUsageCache(cache)) + session_id: Final = "mixed-rollout-session" + legacy_key: Final = handler._make_legacy_cache_key(session_id) + await cache.async_set_cache(key=legacy_key, value=3) + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(20, "agent-a")) + registry.register_agent(_make_mock_agent(20, "agent-b")) + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + assert await handler._increment_agent_and_get(session_id, "agent-a") == 4 + assert await handler._increment_agent_and_get(session_id, "agent-b") == 4 + + # An old pod still writes only the legacy aggregate. + await cache.async_set_cache(key=legacy_key, value=6) + assert await handler._increment_agent_and_get(session_id, "agent-a") == 6 + assert await handler._increment_agent_and_get(session_id, "agent-b") == 6 + + +@pytest.mark.asyncio +async def test_agent_iteration_redis_error_is_not_retried_or_counted_locally() -> None: + handler: Final = _PROXY_MaxIterationsHandler(InternalUsageCache(DualCache())) + calls = 0 + + async def fail_after_attempt(**_kwargs: object) -> object: + nonlocal calls + calls += 1 + raise TimeoutError("reply timed out after Redis may have applied the script") + + handler.increment_script = fail_after_attempt + with patch.object(handler, "_get_local_count", new=AsyncMock(side_effect=AssertionError("must fail closed"))): + with pytest.raises(TimeoutError): + await handler._increment_agent_and_get("uncertain-session", "agent-a") + + assert calls == 1 + + +@pytest.mark.asyncio +@pytest.mark.skipif(_redis_port_for_migration_test() is None, reason="requires local Redis for Lua migration path") +async def test_redis_iteration_migration_is_atomic_and_preserves_session_ttl() -> None: + from litellm.caching.redis_cache import RedisCache + + port: Final = _redis_port_for_migration_test() + assert port is not None + redis: Final = RedisCache(host="127.0.0.1", port=port) + redis_client: Final = redis.init_async_client() + handler: Final = _PROXY_MaxIterationsHandler(InternalUsageCache(DualCache(redis_cache=redis))) + handler.ttl = 30 + registry: Final = AgentRegistry() + registry.register_agent(_make_mock_agent(200, "agent-a")) + registry.register_agent(_make_mock_agent(200, "agent-b")) + session_id: Final = f"iterations-migration-{uuid.uuid4().hex}" + keys: Final = [] + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry): + try: + legacy_key: Final = handler._make_legacy_cache_key(session_id) + total_new_key: Final = handler._make_total_new_cache_key(session_id) + keys.extend( + ( + legacy_key, + total_new_key, + handler._make_cache_key(session_id, "agent-a"), + handler._make_cache_key(session_id, "agent-b"), + ) + ) + await redis.async_set_cache(key=legacy_key, value=3, ttl=30) + legacy_redis_key: Final = redis.check_and_fix_namespace(legacy_key) + initial_ttl: Final = await redis_client.pttl(legacy_redis_key) + hash_tags: Final = {key[key.index("{") : key.index("}") + 1] for key in keys} + assert len(hash_tags) == 1 + + assert await handler._increment_agent_and_get(session_id, "agent-a") == 4 + assert await handler._increment_agent_and_get(session_id, "agent-b") == 4 + await redis_client.incr(legacy_redis_key) # Old pod updates only the aggregate. + assert await handler._increment_agent_and_get(session_id, "agent-a") == 6 + assert await handler._increment_agent_and_get(session_id, "agent-b") == 6 + + sidecar_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(total_new_key)) + agent_ttl: Final = await redis_client.pttl( + redis.check_and_fix_namespace(handler._make_cache_key(session_id, "agent-a")) + ) + current_ttl: Final = await redis_client.pttl(legacy_redis_key) + assert current_ttl <= sidecar_ttl <= current_ttl + 1100 + assert sidecar_ttl <= agent_ttl <= sidecar_ttl + 1100 + assert current_ttl <= initial_ttl + await asyncio.sleep(0.25) + before_increment_ttl: Final = await redis_client.pttl(legacy_redis_key) + before_increment_sidecar_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(total_new_key)) + assert await handler._increment_agent_and_get(session_id, "agent-a") == 7 + after_increment_ttl: Final = await redis_client.pttl(legacy_redis_key) + after_increment_sidecar_ttl: Final = await redis_client.pttl(redis.check_and_fix_namespace(total_new_key)) + assert after_increment_ttl <= before_increment_ttl + 50 + assert after_increment_sidecar_ttl <= before_increment_sidecar_ttl + 50 + + values: Final = await asyncio.gather( + *(handler._increment_agent_and_get(session_id, "agent-a") for _ in range(100)) + ) + assert sorted(values) == list(range(8, 108)) + + await redis_client.pexpire(legacy_redis_key, 100) + await asyncio.sleep(0.15) + with pytest.raises(Exception, match="legacy session count expired before agent scope"): + await handler._increment_agent_and_get(session_id, "agent-a") + + await redis_client.delete(redis.check_and_fix_namespace(total_new_key)) + with pytest.raises(Exception, match="agent session count exists without migration total"): + await handler._increment_agent_and_get(session_id, "agent-a") + finally: + await redis_client.delete(*(redis.check_and_fix_namespace(key) for key in keys)) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py index 879f2b1d889..556a3f1068d 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py @@ -953,7 +953,7 @@ async def test_max_budget_per_session_limiter_populates_provider(): ) mock_registry.stable_agent_id.return_value = "agent-session-budget" with patch.object( - handler, "_get_current_spend", new=AsyncMock(return_value=5.0) + handler, "_get_agent_spend", new=AsyncMock(return_value=5.0) ): with pytest.raises(HTTPException) as exc_info: await handler.async_pre_call_hook( @@ -989,7 +989,7 @@ async def test_max_budget_per_session_limiter_unknown_model_falls_back(): ) mock_registry.stable_agent_id.return_value = "agent-session-budget" with patch.object( - handler, "_get_current_spend", new=AsyncMock(return_value=5.0) + handler, "_get_agent_spend", new=AsyncMock(return_value=5.0) ): with pytest.raises(HTTPException) as exc_info: await handler.async_pre_call_hook( diff --git a/tests/unit/test_rate_limit_error_unification.py b/tests/unit/test_rate_limit_error_unification.py index 0e0e5273e6d..abef8cf512e 100644 --- a/tests/unit/test_rate_limit_error_unification.py +++ b/tests/unit/test_rate_limit_error_unification.py @@ -841,6 +841,7 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError: ) internal_cache = MagicMock() + internal_cache.dual_cache = DualCache() internal_cache.async_get_cache = AsyncMock(return_value=10.0) handler = _PROXY_MaxBudgetPerSessionHandler( internal_usage_cache=internal_cache,