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:
user
2026-04-23 02:48:03 +00:00
parent e0a9c193f8
commit fb55e2f1b9
2 changed files with 62 additions and 7 deletions
+9 -7
View File
@@ -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