scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): enforce user budget on team keys and fail closed in max_budget_limiter

Browse source 
Align _PROXY_MaxBudgetLimiter with common_checks and budget reservation: only exempt team keys when skip_user_budget_on_team_key is set, and let spend-lookup errors propagate instead of swallowing them (fail closed). Fixes #33323
This commit is contained in:
Devin AI 2026-07-16 15:40:48 +00:00
parent 69a491e168
commit 98960f80f2
2 changed files with 129 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

103
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -1,6 +1,3 @@
from fastapi import HTTPException
from litellm import verbose_logger
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
@ -22,60 +19,54 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
data: dict,
call_type: str,
):
try:
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
max_budget = user_api_key_dict.user_max_budget
user_id = user_api_key_dict.user_id
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
max_budget = user_api_key_dict.user_max_budget
user_id = user_api_key_dict.user_id
if max_budget is None or user_id is None:
return
if max_budget is None or user_id is None:
return
# Personal budget applies only to non-team requests, matching
# the explicit team-key exemption in common_checks section 4.1.
if user_api_key_dict.team_id is not None:
return
from litellm.proxy.proxy_server import general_settings
# The reservation path admits at the strict-`<` boundary and
# atomically pre-fills the same counter we'd read here. Re-checking
# with `>=` would reject a request the reservation already admitted
# when the reservation fills the counter to exactly max_budget.
# Imported lazily to avoid a circular import via proxy.utils.
from litellm.proxy.spend_tracking.budget_reservation import (
get_reserved_counter_keys,
)
user_counter_key = f"spend:user:{user_id}"
if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation):
return
from litellm.proxy.proxy_server import get_current_spend
curr_spend = await get_current_spend(
counter_key=user_counter_key,
fallback_spend=user_api_key_dict.user_spend or 0.0,
)
verbose_proxy_logger.debug(
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
user_id,
curr_spend,
max_budget,
)
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None)
raise ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)
except HTTPException as e:
raise e
except Exception as e:
verbose_logger.exception(
"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - {}".format(
str(e)
)
skip_for_team = (
general_settings.get("skip_user_budget_on_team_key") is True and user_api_key_dict.team_id is not None
)
if skip_for_team:
return
# The reservation path admits at the strict-`<` boundary and
# atomically pre-fills the same counter we'd read here. Re-checking
# with `>=` would reject a request the reservation already admitted
# when the reservation fills the counter to exactly max_budget.
# Imported lazily to avoid a circular import via proxy.utils.
from litellm.proxy.spend_tracking.budget_reservation import (
get_reserved_counter_keys,
)
user_counter_key = f"spend:user:{user_id}"
if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation):
return
from litellm.proxy.proxy_server import get_current_spend
curr_spend = await get_current_spend(
counter_key=user_counter_key,
fallback_spend=user_api_key_dict.user_spend or 0.0,
)
verbose_proxy_logger.debug(
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
user_id,
curr_spend,
max_budget,
)
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None)
raise ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)

84
tests/test_litellm/proxy/hooks/test_max_budget_limiter.py
View file

@ -163,14 +163,47 @@ async def test_does_not_skip_when_reservation_covers_a_different_counter():
@pytest.mark.asyncio
async def test_team_keys_skip_personal_budget():
async def test_team_keys_enforce_personal_budget_by_default():
"""With the default policy (`skip_user_budget_on_team_key` unset/false) a
team key must still enforce the user's personal budget, matching
`common_checks` and budget reservation. Regression for #33323."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch(
with patch.dict("litellm.proxy.proxy_server.general_settings", {}, clear=True), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=20.0),
) as mock_get_spend:
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
mock_get_spend.assert_awaited()
@pytest.mark.asyncio
async def test_team_keys_skip_personal_budget_when_flag_set():
"""`skip_user_budget_on_team_key=True` restores the legacy behavior where a
team key skips the user's personal budget."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"skip_user_budget_on_team_key": True},
clear=True,
), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
) as mock_get_spend:
@ -185,6 +218,53 @@ async def test_team_keys_skip_personal_budget():
mock_get_spend.assert_not_awaited()
@pytest.mark.asyncio
async def test_non_team_key_unaffected_by_skip_flag():
"""The `skip_user_budget_on_team_key` flag only exempts team keys; a
non-team key must still be enforced even when the flag is set."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0, team_id=None)
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"skip_user_budget_on_team_key": True},
clear=True,
), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=20.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_spend_lookup_failure_fails_closed():
"""A spend-lookup infrastructure error (e.g. cache down) must propagate so
the hook fails closed instead of silently admitting the request.
Regression for #33323."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0, team_id=None)
with patch.dict("litellm.proxy.proxy_server.general_settings", {}, clear=True), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(side_effect=RuntimeError("cache down")),
):
with pytest.raises(RuntimeError, match="cache down"):
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
@pytest.mark.asyncio
async def test_no_max_budget_passes():
handler = _PROXY_MaxBudgetLimiter()