mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-17 16:17:18 +00:00
fix(auth): propagate BudgetExceededError through _safe_fetch
get_end_user_object raises litellm.BudgetExceededError internally when the end-user is over budget. The previous _safe_fetch in the centralized gate swallowed it and returned None, which caused common_checks to see end_user_object=None and skip the budget check entirely — silently bypassing end-user budget enforcement. Add BudgetExceededError to the re-raise list alongside HTTPException and ProxyException (reported by Veria AI).
This commit is contained in:
@@ -1525,16 +1525,18 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
||||
|
||||
|
||||
async def _safe_fetch(label: str, awaitable):
|
||||
"""Run an awaitable and return its result, logging and swallowing
|
||||
HTTPException / ProxyException / PrismaError. Used in the centralized
|
||||
authz gate so a DB outage fetching a team/user/project yields a
|
||||
``None`` object (and ``common_checks`` no-ops the corresponding
|
||||
branch) rather than failing the request before authorization can
|
||||
run against whatever limits are recorded on the token itself.
|
||||
"""Run an awaitable and return its result. Re-raises authentication /
|
||||
authorization failures (HTTPException, ProxyException,
|
||||
BudgetExceededError — which ``get_end_user_object`` raises for
|
||||
end-user budget violations) so they propagate to the caller.
|
||||
Other exceptions (e.g. transient DB errors fetching context) are
|
||||
swallowed with a debug log and ``None`` is returned so
|
||||
``common_checks`` can still run against whatever limits are recorded
|
||||
directly on the token.
|
||||
"""
|
||||
try:
|
||||
return await awaitable
|
||||
except (HTTPException, ProxyException) as e:
|
||||
except (HTTPException, ProxyException, litellm.BudgetExceededError) as e:
|
||||
verbose_proxy_logger.debug(
|
||||
"centralized auth: %s fetch failed (%s: %s)",
|
||||
label,
|
||||
|
||||
@@ -1960,6 +1960,59 @@ async def test_centralized_common_checks_tolerates_db_errors_when_fetching_conte
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_propagates_end_user_budget_error():
|
||||
"""Regression: ``get_end_user_object`` raises ``litellm.BudgetExceededError``
|
||||
internally when an end user is over budget. ``_safe_fetch`` must
|
||||
re-raise it so the wrapper surfaces the budget violation, rather
|
||||
than swallowing it and letting ``common_checks`` see
|
||||
``end_user_object=None`` and skip enforcement."""
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
token = UserAPIKeyAuth(api_key="sk-test", user_id="u", end_user_id="alice")
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
request._body = json.dumps({"user": "alice", "model": "gpt-4o"}).encode()
|
||||
|
||||
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_end_user_object",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=litellm.BudgetExceededError(
|
||||
message="End-user budget exceeded",
|
||||
current_cost=20.0,
|
||||
max_budget=10.0,
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_checks,
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await _run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data={"user": "alice", "model": "gpt-4o"},
|
||||
route="/chat/completions",
|
||||
)
|
||||
# common_checks must not be invoked if the budget violation
|
||||
# propagates from the context gathering — the wrapper should
|
||||
# fail before reaching it.
|
||||
mock_checks.assert_not_awaited()
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_short_circuits_when_master_key_unset():
|
||||
"""master_key=None is no-auth dev mode — admin-only routes and
|
||||
|
||||
Reference in New Issue
Block a user