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

@@ -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
+    )

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

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()
+    )