From 85b59205287b061b6ac85550ca3eb44641eae9e4 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 5 Jun 2026 12:09:47 -0700 Subject: [PATCH] feat(proxy): auth_v2 enforces team/org/global budgets by reusing v1's functions Closes the hierarchy-budget gap: team, organization, and global caps live in v1's common_checks (not the pre-call hooks), which auth_v2 does not run, so they were unenforced under v2. enforce_hierarchy_budgets calls the exact same functions v1 uses - _team_max_budget_check, _organization_max_budget_check, and get_global_proxy_spend + _global_proxy_budget_check - so there is one budget implementation with two callers (true single authority), with the correct spend-counter conventions and no edit to v1's path. Wired into the inference branch for all login types; a breach surfaces as the same 429 ProxyException v1 raises. Verified in-process: a team over budget is blocked (BudgetExceededError), under budget is allowed, and a teamless identity is a no-op. Real-counter accuracy across pods remains a live check (see INTEGRATION.md). --- litellm/proxy/auth/v2/budgets.py | 56 +++++++++++++ litellm/proxy/auth/v2/entry.py | 19 +++++ .../proxy/auth/v2/test_budgets.py | 78 +++++++++++++++++++ 3 files changed, 153 insertions(+) create mode 100644 litellm/proxy/auth/v2/budgets.py create mode 100644 tests/test_litellm/proxy/auth/v2/test_budgets.py diff --git a/litellm/proxy/auth/v2/budgets.py b/litellm/proxy/auth/v2/budgets.py new file mode 100644 index 00000000000..0d91d298bf4 --- /dev/null +++ b/litellm/proxy/auth/v2/budgets.py @@ -0,0 +1,56 @@ +from typing import Any, Optional + + +async def enforce_hierarchy_budgets(identity: Any, route: str, ctx: Any) -> None: + """Enforce team, organization, and global budgets for an auth_v2 request. + + These hierarchy caps live in v1's ``common_checks`` (not the pre-call hooks), + which auth_v2 does not run -- so without this they would go unenforced under + v2. Rather than reimplement them (divergent logic, wrong counter keys), this + calls the exact same functions v1 uses, so there is one budget implementation + with two callers. Raises ``litellm.BudgetExceededError`` when a cap is + exceeded; the key/user budgets are already handled by the pre-call hooks. + """ + from litellm.proxy.auth.auth_checks import ( + _global_proxy_budget_check, + _organization_max_budget_check, + _team_max_budget_check, + get_team_object, + ) + from litellm.proxy.auth.user_api_key_auth import get_global_proxy_spend + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + team_id = getattr(identity, "team_id", None) + team_object: Optional[Any] = None + if team_id is not None: + team_object = await get_team_object( + team_id=team_id, + prisma_client=ctx.prisma_client, + user_api_key_cache=ctx.user_api_key_cache, + parent_otel_span=ctx.parent_otel_span, + proxy_logging_obj=ctx.proxy_logging_obj, + ) + + await _team_max_budget_check( + team_object=team_object, + valid_token=identity, + proxy_logging_obj=ctx.proxy_logging_obj, + ) + await _organization_max_budget_check( + valid_token=identity, + team_object=team_object, + prisma_client=ctx.prisma_client, + user_api_key_cache=ctx.user_api_key_cache, + proxy_logging_obj=ctx.proxy_logging_obj, + ) + + global_proxy_spend = await get_global_proxy_spend( + litellm_proxy_admin_name=litellm_proxy_admin_name, + user_api_key_cache=ctx.user_api_key_cache, + prisma_client=ctx.prisma_client, + token=getattr(identity, "token", None) or "", + proxy_logging_obj=ctx.proxy_logging_obj, + ) + _global_proxy_budget_check( + global_proxy_spend=global_proxy_spend, skip_budget_checks=False, route=route + ) diff --git a/litellm/proxy/auth/v2/entry.py b/litellm/proxy/auth/v2/entry.py index 62d5248dc8a..1bef310ad80 100644 --- a/litellm/proxy/auth/v2/entry.py +++ b/litellm/proxy/auth/v2/entry.py @@ -6,6 +6,7 @@ from litellm.integrations.otel.runtime import seed_request_identity from .authenticators import AuthContext, AuthResult, authenticate from .authorizer import AuthorizationDenied, authorize +from .budgets import enforce_hierarchy_budgets from .context import AuthMethod, RequestAuthContext, set_auth_context from .end_user import resolve_end_user from .enforcer import CasbinEnforcer @@ -69,6 +70,23 @@ async def _enrich_for_limits(identity: Any, ctx: AuthContext) -> None: await enrich_identity(identity, load_user=load_user, load_team=load_team) +async def _enforce_budgets(identity: Any, route: str, ctx: AuthContext) -> None: + """Enforce team/org/global budgets (reusing v1's functions) and surface a + breach as the same ProxyException v1 raises.""" + import litellm + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + try: + await enforce_hierarchy_budgets(identity, route, ctx) + except litellm.BudgetExceededError as e: + raise ProxyException( + message=e.message, + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), + ) + + async def _best_effort_identity(api_key: Optional[str], ctx: AuthContext) -> AuthResult: """On loud-open routes, use the real identity if a usable key is present, otherwise fall back to an anonymous principal. Never fails the request.""" @@ -160,6 +178,7 @@ async def user_api_key_auth_v2( status_code=status.HTTP_403_FORBIDDEN, detail=f"auth_v2: not permitted to call model '{requested_model}'", ) + await _enforce_budgets(identity, route, ctx) await resolve_end_user(request, request_data, dict(request.headers)) seed_request_identity(identity, model=requested_model) identity.request_route = route diff --git a/tests/test_litellm/proxy/auth/v2/test_budgets.py b/tests/test_litellm/proxy/auth/v2/test_budgets.py new file mode 100644 index 00000000000..f01237cc758 --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_budgets.py @@ -0,0 +1,78 @@ +"""Verify auth_v2 enforces team budgets by reusing v1's exact budget functions. + +enforce_hierarchy_budgets calls the same _team_max_budget_check v1 uses in +common_checks (single authority). Drives it with a team over/under budget, +mocking only get_team_object and the spend counter; the global cap is inert here +(litellm.max_budget defaults to 0). Real-counter accuracy is the live piece. +""" + +from types import SimpleNamespace + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.v2.budgets import enforce_hierarchy_budgets + + +class _Logging: + async def budget_alerts(self, **kwargs): + return None + + +def _ctx(): + return SimpleNamespace( + prisma_client=None, + user_api_key_cache=None, + proxy_logging_obj=_Logging(), + parent_otel_span=None, + ) + + +@pytest.fixture +def team(monkeypatch): + def _set(max_budget, current_spend): + team_obj = SimpleNamespace( + team_id="t1", max_budget=max_budget, spend=0.0, organization_id=None + ) + + async def fake_get_team_object(**kwargs): + return team_obj + + async def fake_get_current_spend(counter_key, fallback_spend): + return current_spend + + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_team_object", fake_get_team_object + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", fake_get_current_spend + ) + + return _set + + +@pytest.mark.asyncio +async def test_team_over_budget_is_blocked(team): + team(max_budget=10.0, current_spend=15.0) + with pytest.raises(litellm.BudgetExceededError): + await enforce_hierarchy_budgets( + UserAPIKeyAuth(team_id="t1"), "/chat/completions", _ctx() + ) + + +@pytest.mark.asyncio +async def test_team_under_budget_is_allowed(team): + team(max_budget=10.0, current_spend=5.0) + await enforce_hierarchy_budgets( + UserAPIKeyAuth(team_id="t1"), "/chat/completions", _ctx() + ) + + +@pytest.mark.asyncio +async def test_no_team_means_no_hierarchy_cap(team): + # A keyless/teamless identity has no team to load; over-spend is irrelevant. + team(max_budget=10.0, current_spend=999.0) + await enforce_hierarchy_budgets( + UserAPIKeyAuth(team_id=None), "/chat/completions", _ctx() + )