diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 1983675c5f3..4c023205dbb 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -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, ) diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py index 0074d7062b8..f25e968a606 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -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()