scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

fix(proxy): reserve session budget at admission to close concurrent max_budget_per_session bypass

Browse source
This commit is contained in:
Devin AI 2026-07-27 15:56:01 +00:00
parent 24123269cc
commit 3d7f68bd02
2 changed files with 384 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

265
litellm/proxy/hooks/max_budget_per_session_limiter.py
View file

@ -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

156
tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py
View file

@ -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)