diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5ba5e8fa1aa..7c09c8af75b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2888,6 +2888,12 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + # Whether this request is exempt from budget enforcement because the requested + # model is zero-cost. Decided once in user_api_key_auth() and published here so + # every enforcement point agrees, including the ones that run after auth. + # exclude=True keeps it out of serialised key responses; the default enforces + # the budget, so a path that never computed the exemption fails closed. + skip_budget_checks: bool = Field(default=False, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e92d090a2fb..5a52dd2e547 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1561,6 +1561,7 @@ async def _user_api_key_auth_builder( skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) + valid_token.skip_budget_checks = skip_budget_checks # Fetch project object for JWT path if project_id is set _jwt_project_obj = None @@ -1984,6 +1985,7 @@ async def _user_api_key_auth_builder( skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) + valid_token.skip_budget_checks = skip_budget_checks # Check 3. Check if user is in their team budget if not skip_budget_checks and valid_token.team_member_spend is not None: @@ -2615,6 +2617,7 @@ async def _run_centralized_common_checks( request=request, llm_router=llm_router, ) + user_api_key_auth_obj.skip_budget_checks = skip_budget_checks # Pin the metadata variable name (litellm_metadata vs metadata) before # any tag merge runs. Without this, header tags from diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index eaf37b0bcf1..d5f3f296f18 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -32,6 +32,19 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): if max_budget is None or user_id is None: return + # A zero-cost model is exempt from the personal budget, as it already is + # from the key, team and end-user budgets. The exemption is not recomputed + # here: user_api_key_auth() decided it and published it on the auth object, + # the same way budget_reservation below is handed in, so this hook cannot + # disagree with the checks that ran during auth. + if user_api_key_dict.skip_budget_checks: + verbose_proxy_logger.debug( + "MaxBudgetLimiter: user_id=%s is over budget but the request is exempt " + "(zero-cost model) - allowing", + user_id, + ) + return + from litellm.proxy.proxy_server import general_settings if ( 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 71671966d1a..e49fa462b53 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -27,6 +27,7 @@ def _make_user_api_key_auth( user_spend: float = 0.0, team_id=None, budget_reservation=None, + skip_budget_checks: bool = False, ) -> UserAPIKeyAuth: return UserAPIKeyAuth( api_key="sk-test", @@ -35,6 +36,7 @@ def _make_user_api_key_auth( user_spend=user_spend, team_id=team_id, budget_reservation=budget_reservation, + skip_budget_checks=skip_budget_checks, ) @@ -235,3 +237,43 @@ async def test_no_max_budget_passes(): assert result is None mock_get_spend.assert_not_awaited() + + +# `get_current_spend` falls back to the caller-supplied `fallback_spend`, which the hook +# passes from `user_api_key_dict.user_spend`, so an over-budget request can be set up +# through the auth object alone -- no need to patch the SDK's own spend lookup. The pair +# below differs only in `skip_budget_checks`, which is what pins the exemption. + + +@pytest.mark.asyncio +async def test_budget_exempt_request_passes_when_over_budget(): + """A request auth marked exempt (zero-cost model) is admitted despite being over budget.""" + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0, user_spend=99.0, skip_budget_checks=True) + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "free-model"}, + call_type="completion", + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_non_exempt_request_still_blocked_when_over_budget(): + """The same request without the exemption is still refused, so the flag is what admits it.""" + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0, user_spend=99.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={"model": "paid-model"}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + assert "Max budget limit reached." in exc_info.value.detail