fix(proxy): make per-model budgets track spend, enforce, and report the same counter

litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py

@@ -56,7 +56,7 @@
 # so the summary's spend is attributed to the same scopes. The list mirrors the
 # fields populated by
 # ``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata``.
-# ``user_api_key_model_max_budget`` / ``user_api_key_end_user_model_max_budget``
+# The three ``*_model_max_budget`` fields
 # are what ``_PROXY_VirtualKeyModelMaxBudgetLimiter`` reads post-call to update
 # the per-model spend caches, so without them the summary spend would never
 # count against the caller's model budget. ``user_api_key_end_user_id`` /
@@ -76,6 +76,7 @@
     "user_api_key_end_user_id",
     "user_api_end_user_max_budget",
     "user_api_key_model_max_budget",
+    "user_api_key_user_model_max_budget",
     "user_api_key_end_user_model_max_budget",
     "litellm_call_id",
     "litellm_parent_otel_span",
@@ -317,10 +318,14 @@ async def _check_summary_model_budget(
     The summary subrequest never passes back through ``user_api_key_auth``, so
     without this gate a caller whose ``model_max_budget`` for
     ``context_management_summary_model`` is exhausted could keep consuming that
-    model via compaction. Mirrors the ``model_max_budget`` /
-    ``end_user_model_max_budget`` enforcement that ``user_api_key_auth`` runs for
-    the client-requested model. Returns True outside the proxy or when no
+    model via compaction. Mirrors the per-model budget enforcement that
+    ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no
     per-model budget is configured.
+
+    All three scopes are checked because the summary's spend is charged to all
+    three: this file propagates the key, user and end-user budgets into the
+    subrequest's metadata, so enforcing only two of them would let compaction
+    increment a counter it can never be refused by.
     """
     if user_api_key_auth is None:
         return True
@@ -347,6 +352,25 @@ async def _check_summary_model_budget(
             )
             return False
 
+    user_model_max_budget: Final = getattr(user_api_key_auth, "user_model_max_budget", None)
+    user_id: Final = getattr(user_api_key_auth, "user_id", None)
+    if isinstance(user_model_max_budget, dict) and user_model_max_budget and user_id is not None:
+        try:
+            await model_max_budget_limiter.is_user_within_model_budget(
+                user_id=user_id,
+                user_model_max_budget=user_model_max_budget,
+                model=summary_model,
+            )
+        except litellm.BudgetExceededError:
+            return False
+        except Exception as e:  # noqa: BLE001  # a budget gate denies on any failure, as the key and end-user scopes do
+            verbose_logger.warning(
+                "compact_20260112: unexpected error during user model-budget check for summary_model=%s; denying: %s",
+                summary_model,
+                e,
+            )
+            return False
+
     end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None)
     end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None)
     if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None:

litellm/proxy/_types.py

@@ -2780,6 +2780,10 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView):  # the expected response ob
     user_email: str | None = None
     user_spend: float | None = None
     user_max_budget: float | None = None
+    # Values stay `object` rather than BudgetConfig: this is the raw JSON column,
+    # and validating it here would make one malformed row fail auth outright.
+    # resolve_model_budget validates the single entry a request actually needs.
+    user_model_max_budget: dict[str, object] | None = None
     request_route: str | None = None
     is_session_token: bool = False
     # Server-only marker set exclusively by the MCP gateway admission path
@@ -2957,6 +2961,8 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase):
     sso_user_id: str | None = None
     teams: list[str] = []  # Just team IDs, not full team objects
     object_permission: LiteLLM_ObjectPermissionTable | None = None
+    model_max_budget: dict | None = None
+    model_max_budget_usage: dict | None = None
 
 
 from litellm.models.config import LiteLLM_Config as LiteLLM_Config  # noqa: E402

litellm/proxy/auth/auth_utils.py

@@ -1795,7 +1795,7 @@ def _format_model_candidates(
     return candidates
 
 
-def _request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool:
+def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool:
     """Whether FastAPI resolved this request to a user-defined pass-through handler.
 
     Reads the marker set by ``create_pass_through_route`` off the dispatched endpoint
@@ -1836,7 +1836,7 @@ def get_model_from_request(
     and does not carry the marker. Built-in provider passthrough routes
     (``/vertex_ai``, ``/gemini``, ...) are separate handlers and keep model enforcement.
     """
-    if _request_dispatched_to_pass_through_endpoint(request):
+    if request_dispatched_to_pass_through_endpoint(request):
         return None
 
     candidates: Final = _extract_model_candidates_from_request(

litellm/proxy/auth/user_api_key_auth.py

@@ -11,6 +11,7 @@
 import fnmatch
 import re
 import secrets
+from collections.abc import Mapping
 from datetime import datetime, timezone
 from typing import Any, Final, NamedTuple, Protocol, Union, cast
 
@@ -186,6 +187,62 @@ async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, mo
     async def get_fallback_model_within_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> str | None: ...
 
 
+class _UserModelBudgetLimiter(Protocol):
+    async def is_user_within_model_budget(
+        self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
+    ) -> bool: ...
+
+
+async def _read_user_model_max_budget(
+    user_id: str | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    parent_otel_span: object,
+    proxy_logging_obj: ProxyLogging,
+) -> dict | None:
+    """The user row's `model_max_budget`, or None when the row cannot be read.
+
+    A user whose row is missing must not be refused: this is a budget lookup,
+    and the main auth path likewise treats an unreadable user as no user.
+    """
+    if user_id is None or prisma_client is None:
+        return None
+    try:
+        user_obj: Final = await get_user_object(
+            user_id=user_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            user_id_upsert=False,
+            parent_otel_span=parent_otel_span,  # pyright: ignore[reportArgumentType]  # Span is a runtime union, not usable in an annotation here
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except Exception as e:  # noqa: BLE001  # mirrors the main path's tolerance
+        verbose_logger.debug("Unable to read user for the per-model budget check: %s", e)
+        return None
+    return getattr(user_obj, "model_max_budget", None)
+
+
+async def _check_user_model_budget(
+    valid_token: UserAPIKeyAuth,
+    model_max_budget_limiter: _UserModelBudgetLimiter,
+    models: list[str],
+) -> None:
+    """Enforce the internal user's own `model_max_budget` across the request's models.
+
+    Separate from the key check: a user's per-model budget caps every key they
+    own, so a caller cannot escape it by minting another key.
+    """
+    user_model_max_budget: Final = valid_token.user_model_max_budget
+    if valid_token.user_id is None or not isinstance(user_model_max_budget, Mapping) or not user_model_max_budget:
+        return
+    for model_name in models:
+        await model_max_budget_limiter.is_user_within_model_budget(
+            user_id=valid_token.user_id,
+            user_model_max_budget=user_model_max_budget,
+            model=model_name,
+        )
+
+
 async def _check_key_model_budget_with_fallback(
     valid_token: UserAPIKeyAuth,
     model_max_budget_limiter: _KeyModelBudgetLimiter,
@@ -1390,6 +1447,7 @@ async def _user_api_key_auth_builder(
                         end_user_id=end_user_id,
                         user_tpm_limit=(user_object.tpm_limit if user_object is not None else None),
                         user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
+                        user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
                         team_member_rpm_limit=(
                             team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
                         ),
@@ -1427,6 +1485,13 @@ async def _user_api_key_auth_builder(
                         if auto_registered is not None:
                             auto_registered.jwt_claims = jwt_claims
                             auto_registered.user_email = user_email
+                            # The auto-registered token is built from the new key's
+                            # columns, which carry no user budget. Carry over the
+                            # already-loaded user row rather than re-reading it, or
+                            # the budget check below has nothing to enforce.
+                            auto_registered.user_model_max_budget = (
+                                user_object.model_max_budget if user_object is not None else None
+                            )
                             valid_token = auto_registered
                             api_key = valid_token.token or ""
 
@@ -1458,6 +1523,28 @@ async def _user_api_key_auth_builder(
                             valid_token.project_metadata = _jwt_project_obj.metadata
                             valid_token.project_alias = _jwt_project_obj.project_alias
 
+                    # JWT auth returns here rather than falling through to the
+                    # virtual-key checks below, so the user's per-model budget
+                    # has to be enforced on this path too. Without it the
+                    # post-call increment still charges the counter and nothing
+                    # ever reads it, which is worse than not tracking at all.
+                    # Guarded by the same flag the virtual-key path uses, or a
+                    # zero-cost model would be refused here and allowed there,
+                    # while the log above claims all budget checks were skipped.
+                    if not skip_budget_checks:
+                        await _check_user_model_budget(
+                            valid_token=cast(UserAPIKeyAuth, valid_token),
+                            model_max_budget_limiter=model_max_budget_limiter,
+                            models=_get_model_names_for_budget_checks(
+                                model=_get_model_from_request_context(
+                                    request_data=request_data,
+                                    route=route,
+                                    request=request,
+                                    llm_router=llm_router,
+                                )
+                            ),
+                        )
+
                     return cast(UserAPIKeyAuth, valid_token)
 
         #### ELSE ####
@@ -1811,6 +1898,12 @@ async def _user_api_key_auth_builder(
                     )
                     user_obj = None
 
+                if user_obj is not None:
+                    # The joint verification-token view carries the key's columns only, so the
+                    # user's own per-model budget reaches enforcement and the post-call
+                    # increment through the row fetched here.
+                    valid_token.user_model_max_budget = user_obj.model_max_budget
+
                 if (
                     user_obj is not None
                     and isinstance(user_obj.metadata, dict)
@@ -1974,6 +2067,14 @@ async def _user_api_key_auth_builder(
                         )
                         current_models = _get_model_names_for_budget_checks(model=current_model)
 
+                    # Check 5a. Internal user model_max_budget
+                    if current_models:
+                        await _check_user_model_budget(
+                            valid_token=valid_token,
+                            model_max_budget_limiter=model_max_budget_limiter,
+                            models=current_models,
+                        )
+
                     # Check 5b. End-user model max budget
                     end_user_mmb: Final = valid_token.end_user_model_max_budget
                     if (
@@ -2757,6 +2858,7 @@ async def _return_user_api_key_auth_obj(
             user_email=user_obj.user_email,
             user_spend=getattr(user_obj, "spend", None),
             user_max_budget=getattr(user_obj, "max_budget", None),
+            user_model_max_budget=getattr(user_obj, "model_max_budget", None),
         )
     if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj):
         user_api_key_kwargs.update(
@@ -3020,10 +3122,21 @@ async def _run_post_custom_auth_checks(
     )
     current_models = _get_model_names_for_budget_checks(model=current_model)
 
+    # A zero-cost model cannot move any counter, so refusing it means refusing on
+    # spend some other model accrued. The JWT and virtual-key paths already skip
+    # every budget check for these; this path did not, so the same request could
+    # be refused under custom auth and served under the other two.
+    skip_budget_checks: Final = (
+        _is_model_cost_zero(model=current_model, llm_router=llm_router)
+        if current_model is not None and llm_router is not None
+        else False
+    )
+
     # 3. Check key-level model_max_budget
     max_budget_per_model: Final = valid_token.model_max_budget
     if (
-        max_budget_per_model is not None
+        not skip_budget_checks
+        and max_budget_per_model is not None
         and isinstance(max_budget_per_model, dict)
         and len(max_budget_per_model) > 0
         and current_models
@@ -3050,10 +3163,33 @@ async def _run_post_custom_auth_checks(
         )
         current_models = _get_model_names_for_budget_checks(model=current_model)
 
+    # 3b. Attach and check the internal user's model_max_budget.
+    # Custom auth builds its own token, so unlike the main path nothing has
+    # loaded the user row yet. The attach is unconditional because the post-call
+    # spend hook reads this field off the token: gating it on the same condition
+    # as enforcement would leave the user's counter uncharged whenever this
+    # request was not itself enforceable, which is the untracked-spend bug this
+    # PR exists to fix.
+    user_budget: Final = await _read_user_model_max_budget(
+        user_id=valid_token.user_id,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        parent_otel_span=parent_otel_span,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    valid_token.user_model_max_budget = user_budget  # rebind-ok: the spend hook reads it off this token
+    if not skip_budget_checks and current_models:
+        await _check_user_model_budget(
+            valid_token=valid_token,
+            model_max_budget_limiter=model_max_budget_limiter,
+            models=current_models,
+        )
+
     # 4. Check end-user model_max_budget
     end_user_mmb: Final = valid_token.end_user_model_max_budget
     if (
-        end_user_mmb is not None
+        not skip_budget_checks
+        and end_user_mmb is not None
         and isinstance(end_user_mmb, dict)
         and len(end_user_mmb) > 0
         and current_models