diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 53de9e0c5a3..bf9ed6c376c 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -2,16 +2,23 @@ Per-Session Budget Limiter for LiteLLM Proxy. Enforces a dollar-amount cap per session (identified by `session_id` / -`x-litellm-trace-id`). After each successful LLM call the response cost is -accumulated against the session. When the accumulated spend exceeds -`max_budget_per_session` (configured in agent litellm_params), subsequent -requests for that session receive a 429. +`x-litellm-trace-id`). Configured via `max_budget_per_session` in agent +litellm_params. When a session's spend would exceed the cap, requests receive +a 429. + +Admission reserves each request's estimated max cost against the session +counter before the call runs (an atomic Redis INCRBYFLOAT, same idea as the +key/team optimistic reservation in budget_reservation.py), so two concurrent +requests for one session can't both read the same below-budget value and both +slip past the gate before either records its cost. After the call the +reservation is reconciled to the actual response cost; on failure it is +refunded. When the request's cost can't be estimated (model not in the cost +map) the hook falls back to read-time enforcement only. Note: trace-id enforcement (require_trace_id_on_calls_by_agent) is handled separately in auth_checks.py at the agent level, not in this hook. Works across multiple proxy instances via DualCache (in-memory + Redis). -Follows the same pattern as max_iterations_limiter.py. """ import os @@ -26,6 +33,8 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache InternalUsageCache = _InternalUsageCache @@ -53,6 +62,12 @@ return new_val # Default TTL for session budget counters (1 hour) DEFAULT_MAX_BUDGET_PER_SESSION_TTL = 3600 +_RESERVED_COST_KEY = "_litellm_session_budget_reserved_cost" +_RESERVATION_RELEASED_KEY = "_litellm_session_budget_reservation_released" +_RESERVATION_SESSION_KEY = "_litellm_session_budget_session_id" + +_EPSILON = 1e-12 + class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): """ @@ -89,11 +104,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): call_type: str, ) -> Optional[Union[Exception, str, dict]]: """ - Before each LLM call, check if max_budget_per_session is set and - whether accumulated spend exceeds the budget (429 if so). + Reserve this request's estimated max cost against the session counter + and reject (429) if the reservation would exceed max_budget_per_session. + Falls back to a read-time spend check when the cost can't be estimated. """ max_budget = self._get_max_budget_per_session(user_api_key_dict) - session_id = self._get_session_id(data) if max_budget is None or session_id is None: @@ -101,35 +116,96 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): max_budget = float(max_budget) cache_key = self._make_cache_key(session_id) + + reservation_cost = self._estimate_reservation_cost(data=data, call_type=call_type) + if reservation_cost is None or reservation_cost <= 0: + await self._enforce_read_time(data=data, session_id=session_id, cache_key=cache_key, max_budget=max_budget) + return None + + await self._reserve_and_enforce( + data=data, + session_id=session_id, + cache_key=cache_key, + max_budget=max_budget, + reservation_cost=reservation_cost, + ) + return None + + async def _enforce_read_time(self, data: dict, session_id: str, cache_key: str, max_budget: float) -> None: + """Read-time only enforcement, used when the request cost is unknown.""" current_spend = await self._get_current_spend(cache_key) verbose_proxy_logger.debug( - "MaxBudgetPerSessionHandler: session_id=%s, spend=%.4f, max=%.2f", + "MaxBudgetPerSessionHandler: session_id=%s, spend=%.4f, max=%.2f (read-time)", session_id, current_spend, max_budget, ) if current_spend >= max_budget: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) - raise ProxyRateLimitError( - detail=( - f"Session budget exceeded for session {session_id}. " - f"Current spend: ${current_spend:.4f}, " - f"max_budget_per_session: ${max_budget:.2f}." - ), - rate_limit_type=RateLimitType.BUDGET, - model=resolved_model, - llm_provider=llm_provider, + raise self._budget_exceeded_error( + data=data, + session_id=session_id, + current_spend=current_spend, + max_budget=max_budget, ) - return None + async def _reserve_and_enforce( + self, + data: dict, + session_id: str, + cache_key: str, + max_budget: float, + reservation_cost: float, + ) -> None: + """ + Atomically reserve `reservation_cost` against the session counter, then + decide admission from the post-increment value so concurrent requests + serialize on the counter instead of racing a stale read. + """ + new_total = await self._increment_spend(cache_key, reservation_cost) + spend_before = new_total - reservation_cost + + if new_total > max_budget: + remaining_before = max_budget - spend_before + if remaining_before > _EPSILON: + await self._increment_spend(cache_key, remaining_before - reservation_cost) + self._stash_reservation(data=data, session_id=session_id, reserved_cost=remaining_before) + return + + await self._increment_spend(cache_key, -reservation_cost) + raise self._budget_exceeded_error( + data=data, + session_id=session_id, + current_spend=spend_before, + max_budget=max_budget, + ) + + self._stash_reservation(data=data, session_id=session_id, reserved_cost=reservation_cost) + + verbose_proxy_logger.debug( + "MaxBudgetPerSessionHandler: reserved %.6f for session %s, spend=%.4f/%.2f", + reservation_cost, + session_id, + new_total, + max_budget, + ) 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. + Reconcile the admission reservation to the call's actual cost. When no + reservation was made (cost couldn't be estimated at admission), fall + back to incrementing the session spend by the full response cost. """ try: + response_cost = float(kwargs.get("response_cost") or 0.0) + + if await self._reconcile_reservation(container=kwargs, actual_cost=response_cost): + return + + if response_cost <= 0: + return + litellm_params = kwargs.get("litellm_params") or {} metadata = litellm_params.get("metadata") or {} session_id = metadata.get("session_id") @@ -149,16 +225,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return agent_litellm_params = agent.litellm_params or {} - max_budget = agent_litellm_params.get("max_budget_per_session") - if max_budget is None: - return - - response_cost = kwargs.get("response_cost") or 0.0 - if response_cost <= 0: + if agent_litellm_params.get("max_budget_per_session") is None: return cache_key = self._make_cache_key(str(session_id)) - await self._increment_spend(cache_key, float(response_cost)) + await self._increment_spend(cache_key, response_cost) verbose_proxy_logger.debug( "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", @@ -171,6 +242,132 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): str(e), ) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + """Refund the admission reservation when the LLM call fails.""" + await self._reconcile_reservation(container=kwargs, actual_cost=0.0) + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + """ + Refund the reservation for a request rejected after admission but before + the LLM call ran (e.g. a downstream guardrail raised). async_log_failure_event + is a completion-level callback and never fires for these proxy-side + rejections, so the reservation would otherwise stay pinned until its TTL. + """ + await self._reconcile_reservation(container=request_data, actual_cost=0.0) + + async def _reconcile_reservation(self, container: dict, actual_cost: float) -> bool: + """ + Adjust the session counter from the reserved amount to `actual_cost`. + + Returns True when a reservation existed (whether or not it still needed + adjusting), so the success path knows to skip the legacy full-cost + increment. Idempotent across the failure/post-call-failure callbacks via + the released marker. + """ + reserved_cost = self._lookup_reserved_cost(container) + if reserved_cost is None: + return False + if self._reservation_released(container): + return True + + session_id = self._lookup_reservation_session_id(container) + if session_id is not None: + adjustment = actual_cost - reserved_cost + if adjustment != 0: + await self._increment_spend(self._make_cache_key(session_id), adjustment) + + self._mark_reservation_released(container) + return True + + def _estimate_reservation_cost(self, data: dict, call_type: str) -> float | None: + try: + from litellm.proxy.proxy_server import llm_router + from litellm.proxy.spend_tracking.budget_reservation import ( + estimate_request_max_cost, + ) + + return estimate_request_max_cost(request_body=data, route=call_type or "", llm_router=llm_router) + except Exception as e: + verbose_proxy_logger.debug( + "MaxBudgetPerSessionHandler: could not estimate request cost, falling back to read-time check: %s", + str(e), + ) + return None + + def _budget_exceeded_error( + self, + data: dict, + session_id: str, + current_spend: float, + max_budget: float, + ) -> "HTTPException": + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) + return ProxyRateLimitError( + detail=( + f"Session budget exceeded for session {session_id}. " + f"Current spend: ${current_spend:.4f}, " + f"max_budget_per_session: ${max_budget:.2f}." + ), + rate_limit_type=RateLimitType.BUDGET, + model=resolved_model, + llm_provider=llm_provider, + ) + + def _stash_reservation(self, data: dict, session_id: str, reserved_cost: float) -> None: + self._stash_in_metadata_channels(data=data, key=_RESERVED_COST_KEY, value=reserved_cost) + self._stash_in_metadata_channels(data=data, key=_RESERVATION_SESSION_KEY, value=session_id) + + @staticmethod + def _stash_in_metadata_channels(data: dict, key: str, value: Any) -> None: + for channel in ("metadata", "litellm_metadata"): + existing = data.get(channel) + if isinstance(existing, dict): + existing[key] = value + elif channel == "metadata": + data[channel] = {key: value} + + def _lookup_reserved_cost(self, container: Any) -> float | None: + value = self._lookup_stashed_value(container, _RESERVED_COST_KEY) + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + def _lookup_reservation_session_id(self, container: Any) -> str | None: + value = self._lookup_stashed_value(container, _RESERVATION_SESSION_KEY) + return str(value) if value is not None else None + + def _reservation_released(self, container: Any) -> bool: + return self._lookup_stashed_value(container, _RESERVATION_RELEASED_KEY) is True + + def _mark_reservation_released(self, container: Any) -> None: + if isinstance(container, dict): + self._stash_in_metadata_channels(data=container, key=_RESERVATION_RELEASED_KEY, value=True) + + @staticmethod + def _lookup_stashed_value(container: Any, key: str) -> Any: + """Resolve a stashed value from any metadata channel a callback sees.""" + if not isinstance(container, dict): + return None + for channel in ("metadata", "litellm_metadata"): + channel_dict = container.get(channel) + if isinstance(channel_dict, dict) and key in channel_dict: + return channel_dict.get(key) + litellm_params = container.get("litellm_params") + if isinstance(litellm_params, dict): + lp_metadata = litellm_params.get("metadata") + if isinstance(lp_metadata, dict) and key in lp_metadata: + return lp_metadata.get(key) + return None + def _get_session_id(self, data: dict) -> Optional[str]: """Extract session_id from request metadata.""" metadata = data.get("metadata") or {} @@ -247,17 +444,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return await self._in_memory_increment_spend(cache_key, amount) async def _in_memory_increment_spend(self, cache_key: str, amount: float) -> float: - current = await self.internal_usage_cache.async_get_cache( + new_value = await self.internal_usage_cache.async_increment_cache( key=cache_key, + value=amount, litellm_parent_otel_span=None, local_only=True, - ) - new_value = (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 + return float(new_value) if new_value is not None else amount 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 879e2d65c7a..4f06a870a42 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,6 +8,7 @@ Tests that session-scoped budget tracking works correctly: - Requests without agent_id pass through """ +import asyncio from unittest.mock import patch import pytest @@ -163,3 +164,158 @@ async def test_no_agent_id_passes(): call_type="", ) assert result is None + + +def _make_handler() -> _PROXY_MaxBudgetPerSessionHandler: + return _PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + + +async def _pre_call(handler, session_id, reservation_cost, agent, user_api_key_dict): + data = {"metadata": {"session_id": session_id}, "model": "gpt-4o"} + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry, patch.object( + handler, "_estimate_reservation_cost", return_value=reservation_cost + ): + mock_registry.get_agent_by_id.return_value = agent + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=handler.internal_usage_cache.dual_cache, + data=data, + call_type="acompletion", + ) + except HTTPException as e: + result = e + return result, data + + +@pytest.mark.asyncio +async def test_concurrent_requests_do_not_bypass_budget(): + """ + Regression for the admission race: several requests firing concurrently + against a fresh session budget must not all slip past the gate. With + per-request reservation only enough requests to fill the budget are + admitted, and the session counter never exceeds max_budget_per_session. + """ + handler = _make_handler() + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-budget", agent_id="agent-budget-123") + agent = _make_mock_agent(max_budget_per_session=1.0) + session_id = "session-concurrent" + + results = await asyncio.gather( + *[_pre_call(handler, session_id, 0.60, agent, user_api_key_dict) for _ in range(5)] + ) + + admitted = [r for r, _ in results if r is None] + rejected = [r for r, _ in results if isinstance(r, HTTPException)] + + assert len(admitted) == 2 + assert len(rejected) == 3 + assert all(r.status_code == 429 for r in rejected) + + spend = await handler._get_current_spend(handler._make_cache_key(session_id)) + assert spend <= 1.0 + 1e-9 + + +@pytest.mark.asyncio +async def test_first_request_admitted_when_estimate_exceeds_budget(): + """ + A single request whose worst-case estimate exceeds the whole budget is + still admitted (reservation resized to the remaining budget), but pins the + counter at the cap so a concurrent sibling is rejected. + """ + handler = _make_handler() + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-budget", agent_id="agent-budget-123") + agent = _make_mock_agent(max_budget_per_session=1.0) + session_id = "session-big-estimate" + + first, _ = await _pre_call(handler, session_id, 5.0, agent, user_api_key_dict) + assert first is None + + spend = await handler._get_current_spend(handler._make_cache_key(session_id)) + assert spend == pytest.approx(1.0) + + second, _ = await _pre_call(handler, session_id, 5.0, agent, user_api_key_dict) + assert isinstance(second, HTTPException) + assert second.status_code == 429 + + +@pytest.mark.asyncio +async def test_reservation_reconciled_to_actual_cost_on_success(): + """ + After admission the reservation is reconciled down to the actual response + cost, so the reserved worst-case does not permanently inflate the session + spend. + """ + handler = _make_handler() + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-budget", agent_id="agent-budget-123") + agent = _make_mock_agent(max_budget_per_session=5.0) + session_id = "session-reconcile" + + result, data = await _pre_call(handler, session_id, 0.80, agent, user_api_key_dict) + assert result is None + + reserved_spend = await handler._get_current_spend(handler._make_cache_key(session_id)) + assert reserved_spend == pytest.approx(0.80) + + await handler.async_log_success_event( + kwargs={"litellm_params": {"metadata": data["metadata"]}, "response_cost": 0.10}, + response_obj=None, + start_time=None, + end_time=None, + ) + + final_spend = await handler._get_current_spend(handler._make_cache_key(session_id)) + assert final_spend == pytest.approx(0.10) + + +@pytest.mark.asyncio +async def test_reservation_refunded_on_failure(): + """A failed call refunds its reservation so it doesn't consume budget.""" + handler = _make_handler() + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-budget", agent_id="agent-budget-123") + agent = _make_mock_agent(max_budget_per_session=5.0) + session_id = "session-refund" + + result, data = await _pre_call(handler, session_id, 0.80, agent, user_api_key_dict) + assert result is None + assert await handler._get_current_spend(handler._make_cache_key(session_id)) == pytest.approx(0.80) + + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": data["metadata"]}}, + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await handler._get_current_spend(handler._make_cache_key(session_id)) == pytest.approx(0.0) + + +@pytest.mark.asyncio +async def test_reservation_refund_is_idempotent_across_failure_hooks(): + """ + The reservation must be refunded at most once even if both + async_post_call_failure_hook and async_log_failure_event fire. + """ + handler = _make_handler() + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test-key-budget", agent_id="agent-budget-123") + agent = _make_mock_agent(max_budget_per_session=5.0) + session_id = "session-idempotent" + + result, data = await _pre_call(handler, session_id, 0.80, agent, user_api_key_dict) + assert result is None + + await handler.async_post_call_failure_hook( + request_data=data, + original_exception=Exception("boom"), + user_api_key_dict=user_api_key_dict, + ) + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": data["metadata"]}}, + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await handler._get_current_spend(handler._make_cache_key(session_id)) == pytest.approx(0.0)