diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
index a7c462a8fb0..2a87afb5990 100644
--- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
+++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
@@ -56,7 +56,7 @@ from ..result import PolyfillResult
# 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 @@ _PROPAGATED_METADATA_KEYS: Final = (
"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:
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 00cbd13cfdc..e51a3138d2d 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2805,6 +2805,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
@@ -2982,6 +2986,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
diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index ce662ee0374..d04a71535ef 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -1801,7 +1801,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
@@ -1842,7 +1842,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(
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 99592d44f9b..fe4f1ee4ae5 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -11,6 +11,7 @@ import asyncio
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 @@ class _KeyModelBudgetLimiter(Protocol):
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
diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py
index 215969ef899..c5d10b2749b 100644
--- a/litellm/proxy/hooks/model_max_budget_limiter.py
+++ b/litellm/proxy/hooks/model_max_budget_limiter.py
@@ -1,21 +1,253 @@
import json
+import time
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
from typing import Final
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import Span
+from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import (
- BudgetConfig,
- GenericBudgetConfigType,
- StandardLoggingPayload,
-)
+from litellm.types.utils import BudgetConfig, StandardLoggingPayload
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend"
+USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend"
+
+_SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
+ {
+ Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
+ Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX,
+ Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX,
+ }
+)
+
+_LEGACY_REQUEST_MODEL_SCOPES: Final = frozenset({Litellm_EntityType.KEY, Litellm_EntityType.END_USER})
+
+_PROCESS_STARTED_AT: Final = time.monotonic()
+
+_BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType(
+ {
+ Litellm_EntityType.KEY: "virtual_key_budget_start_time",
+ Litellm_EntityType.USER: "user_model_budget_start_time",
+ Litellm_EntityType.END_USER: "end_user_budget_start_time",
+ }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedModelBudget:
+ """The `model_max_budget` entry a request resolved to.
+
+ ``budget_model`` is the key as the operator configured it, not the model
+ name on the request. Every counter is keyed on it so enforcement, the
+ post-call increment and the `/key/info` + `/user/info` usage reads cannot
+ disagree about which counter a request belongs to.
+ """
+
+ budget_model: str
+ budget_config: BudgetConfig
+
+
+def model_budget_spend_cache_key(
+ entity_type: Litellm_EntityType,
+ entity_id: str | None,
+ budget_model: str,
+ budget_duration: str | None,
+) -> str:
+ """Sole owner of the per-model spend counter key, shared by its writer and all of its readers."""
+ return f"{_SPEND_CACHE_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}"
+
+
+def _legacy_request_model_spend_cache_key(
+ entity_type: Litellm_EntityType,
+ entity_id: str | None,
+ model: str,
+ resolved: ResolvedModelBudget,
+) -> str | None:
+ """The counter this request was billed to before the budget model owned the key, or None.
+
+ Upgrading proxies carry live counters keyed on the model as REQUESTED
+ (`openai/gpt-4`) rather than as configured (`gpt-4`), and those were the
+ counters the previous version enforced on. Nothing writes that spelling once
+ this version is running, so the pre-upgrade and post-upgrade counters hold
+ disjoint halves of one window and adding them is the window's real spend.
+
+ Only the key and end-user scopes ever had one. The user scope is introduced
+ by this change, so it has no counter to carry.
+
+ The carry stops one budget window after start-up, because a legacy counter
+ belongs to a window that was already open when this process replaced the one
+ writing it. Past that point the lookup could only ever miss.
+ """
+ budget_duration: Final = resolved.budget_config.budget_duration
+ if entity_type not in _LEGACY_REQUEST_MODEL_SCOPES or budget_duration is None:
+ return None
+ if time.monotonic() - _PROCESS_STARTED_AT >= duration_in_seconds(budget_duration):
+ return None
+ return model_budget_spend_cache_key(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ budget_model=model,
+ budget_duration=budget_duration,
+ )
+
+
+def model_budget_start_time_cache_key(
+ entity_type: Litellm_EntityType,
+ entity_id: str | None,
+ budget_model: str,
+ budget_duration: str | None,
+) -> str:
+ """Window start for one (entity, budget model) pair.
+
+ Scoped per budget model because an entity may budget two models over
+ different periods, and a shared start time lets the shorter period restart
+ the longer one's window.
+ """
+ return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}"
+
+
+def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None:
+ """Find the `model_max_budget` entry that governs `model`, or None."""
+ for candidate in _budget_model_candidates(model):
+ raw_budget_config = model_max_budget.get(candidate)
+ if raw_budget_config is None:
+ continue
+ if (budget_config := _usable_budget_config(raw_budget_config)) is None:
+ # An entry that will not validate cannot be keyed, so it cannot be
+ # enforced or incremented. Skip to the next candidate rather than
+ # raising: raising would abort every other scope's increment and turn
+ # a config typo into a 500, and stopping here would let one malformed
+ # specific entry disable a perfectly good bare-family budget beside
+ # it. The candidate chain already falls through an ABSENT entry, and
+ # an unparseable one is indistinguishable from absent to enforcement.
+ # `validate_model_max_budget` rejects these on the write path, so
+ # reaching here means config.yaml or a direct DB edit.
+ verbose_proxy_logger.warning(
+ "Ignoring unusable model_max_budget entry for %s; it cannot be enforced or tracked",
+ candidate,
+ )
+ continue
+ return ResolvedModelBudget(budget_model=candidate, budget_config=budget_config)
+ return None
+
+
+def _budget_model_candidates(model: str) -> tuple[str, ...]:
+ """Names a budget may be configured under for a request on `model`, most specific first.
+
+ Beyond the model as sent, a budget may be keyed on the model without its
+ ``{custom_llm_provider}/`` prefix (``gpt-4o`` governs ``openai/gpt-4o``), on
+ the Bedrock base model (``anthropic.claude-opus-4-8`` governs the
+ cross-region ``us.anthropic.claude-opus-4-8``), or on the bare family name
+ that Bedrock id shares with its direct-provider twin (``claude-opus-4-8``).
+ """
+ return tuple(dict.fromkeys((model, model.split("/")[-1], *_bedrock_candidates(model))))
+
+
+def _bedrock_candidates(model: str) -> tuple[str, ...]:
+ """Bedrock-only candidates, empty unless litellm prices `model` as a Bedrock model.
+
+ Gating on the cost map rather than on a vendor allowlist is what makes
+ splitting the leading dotted segment safe: most dotted model ids are not
+ Bedrock ids at all (``azure/gpt-4.1``, ``gpt-image-1.5``), and splitting one
+ of those would produce a garbage candidate.
+ """
+ base_model: Final = get_bedrock_base_model(model)
+ cost_entry: Final = litellm.model_cost.get(base_model)
+ if not isinstance(cost_entry, dict) or not str(cost_entry.get("litellm_provider", "")).startswith("bedrock"):
+ return ()
+ _, _, without_vendor = base_model.partition(".")
+ return (base_model, without_vendor) if without_vendor else (base_model,)
+
+
+async def build_model_max_budget_usage(
+ entity_type: Litellm_EntityType,
+ entity_id: str | None,
+ model_max_budget: Mapping[str, object] | None,
+ cache: DualCache | None,
+) -> dict[str, dict[str, object]]:
+ """Current-window spend per configured budget model, as `/key/info` and `/user/info` report it.
+
+ `cache` must be the DualCache the limiter writes the counters to; callers
+ read it off the limiter rather than re-deriving it, so a scope that is being
+ blocked can never report zero usage.
+ """
+ if cache is None or entity_id is None or not model_max_budget:
+ return {}
+
+ budgets: Final = tuple(
+ (budget_model, budget_config)
+ for budget_model, raw_budget_config in model_max_budget.items()
+ for budget_config in (_usable_budget_config(raw_budget_config),)
+ if budget_config is not None
+ )
+ if not budgets:
+ return {}
+ spend_keys: Final = tuple(
+ model_budget_spend_cache_key(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ budget_model=budget_model,
+ budget_duration=budget_config.budget_duration,
+ )
+ for budget_model, budget_config in budgets
+ )
+ batched: Final = await cache.async_batch_get_cache(
+ keys=list(spend_keys) # mutable-ok: async_batch_get_cache annotates keys as list, so one must exist here
+ )
+ # async_batch_get_cache returns None if it fails internally, and its result is
+ # index-aligned with `keys` otherwise. An unusable result reads as a miss,
+ # which is what a never-written counter already reads as.
+ current_spends: Final = (
+ tuple(batched) if isinstance(batched, list) and len(batched) == len(budgets) else (None,) * len(budgets)
+ )
+ return {
+ budget_model: {
+ "current_spend": round(_as_spend(current_spend), 4),
+ "budget_limit": budget_config.max_budget,
+ "time_period": budget_config.budget_duration,
+ }
+ for (budget_model, budget_config), current_spend in zip(budgets, current_spends, strict=True)
+ }
+
+
+def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None:
+ try:
+ budget_config: Final = BudgetConfig.model_validate(raw_budget_config)
+ if budget_config.budget_duration is None:
+ return None
+ duration_in_seconds(budget_config.budget_duration)
+ except Exception: # noqa: BLE001 # a malformed entry must not fail the whole report
+ return None
+ return budget_config
+
+
+def _as_spend(current_spend: object) -> float:
+ try:
+ return float(current_spend or 0.0) # pyright: ignore[reportArgumentType] # non-numeric falls to the except
+ except (TypeError, ValueError):
+ return 0.0
+
+
+def _resolve_entity_model_budgets(
+ model: str,
+ entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]],
+) -> tuple[tuple[Litellm_EntityType, str, ResolvedModelBudget], ...]:
+ """Drop the scopes that do not budget `model`, keeping only what can be incremented."""
+ return tuple(
+ (entity_type, entity_id, resolved)
+ for entity_type, entity_id, model_max_budget in entity_budgets
+ if entity_id is not None and isinstance(model_max_budget, Mapping) and model_max_budget
+ for resolved in (resolve_model_budget(model=model, model_max_budget=model_max_budget),)
+ if resolved is not None and resolved.budget_config.budget_duration is not None
+ )
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
@@ -41,47 +273,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
Raises:
BudgetExceededError: If the user_api_key_dict has exceeded the model budget
"""
- _model_max_budget: Final = user_api_key_dict.model_max_budget
- internal_model_max_budget: Final[GenericBudgetConfigType] = {}
-
- for _model, _budget_info in _model_max_budget.items():
- internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
-
- verbose_proxy_logger.debug(
- "internal_model_max_budget %s",
- json.dumps(internal_model_max_budget, indent=4, default=str),
+ return await self._is_entity_within_model_budget(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id=user_api_key_dict.token,
+ model_max_budget=user_api_key_dict.model_max_budget,
+ model=model,
+ exceeded_message=(
+ f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, "
+ f"exceeded budget for model={model}"
+ ),
)
- # check if current model is in internal_model_max_budget
- _current_model_budget_info: Final = self._get_request_model_budget_config(
- model=model, internal_model_max_budget=internal_model_max_budget
- )
- if _current_model_budget_info is None:
- verbose_proxy_logger.debug("Model %s not found in internal_model_max_budget", model)
- return True
-
- # check if current model is within budget
- if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0:
- _current_spend: Final = await self._get_virtual_key_spend_for_model(
- user_api_key_hash=user_api_key_dict.token,
- model=model,
- key_budget_config=_current_model_budget_info,
- )
- if (
- _current_spend is not None
- and _current_model_budget_info.max_budget is not None
- and _current_spend > _current_model_budget_info.max_budget
- ):
- raise litellm.BudgetExceededError(
- message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
- current_cost=_current_spend,
- max_budget=_current_model_budget_info.max_budget,
- entity_type=Litellm_EntityType.KEY.value,
- entity_id=user_api_key_dict.token,
- )
-
- return True
-
async def get_fallback_model_within_budget(
self,
user_api_key_dict: UserAPIKeyAuth,
@@ -96,10 +298,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
continue
return None
+ async def is_user_within_model_budget(
+ self,
+ user_id: str,
+ user_model_max_budget: Mapping[str, object],
+ model: str,
+ ) -> bool:
+ """
+ Check if the internal user is within the model budget
+
+ Raises:
+ BudgetExceededError: If the user has exceeded the model budget
+ """
+ return await self._is_entity_within_model_budget(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ model_max_budget=user_model_max_budget,
+ model=model,
+ exceeded_message=f"LiteLLM User: {user_id}, exceeded budget for model={model}",
+ )
+
async def is_end_user_within_model_budget(
self,
end_user_id: str,
- end_user_model_max_budget: dict,
+ end_user_model_max_budget: Mapping[str, object],
model: str,
) -> bool:
"""
@@ -108,116 +330,81 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
Raises:
BudgetExceededError: If the end_user has exceeded the model budget
"""
- internal_model_max_budget: Final[GenericBudgetConfigType] = {}
-
- for _model, _budget_info in end_user_model_max_budget.items():
- internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
-
- verbose_proxy_logger.debug(
- "end_user internal_model_max_budget %s",
- json.dumps(internal_model_max_budget, indent=4, default=str),
+ return await self._is_entity_within_model_budget(
+ entity_type=Litellm_EntityType.END_USER,
+ entity_id=end_user_id,
+ model_max_budget=end_user_model_max_budget,
+ model=model,
+ exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}",
)
- # check if current model is in internal_model_max_budget
- _current_model_budget_info: Final = self._get_request_model_budget_config(
- model=model, internal_model_max_budget=internal_model_max_budget
- )
- if _current_model_budget_info is None:
- verbose_proxy_logger.debug("Model %s not found in end_user_model_max_budget", model)
+ async def _is_entity_within_model_budget(
+ self,
+ entity_type: Litellm_EntityType,
+ entity_id: str | None,
+ model_max_budget: Mapping[str, object] | None,
+ model: str,
+ exceeded_message: str,
+ ) -> bool:
+ if not model_max_budget:
+ return True
+ resolved: Final = resolve_model_budget(model=model, model_max_budget=model_max_budget)
+ if resolved is None:
+ verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value)
return True
- # check if current model is within budget
- if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0:
- _current_spend: Final = await self._get_end_user_spend_for_model(
- end_user_id=end_user_id,
- model=model,
- key_budget_config=_current_model_budget_info,
- )
- if (
- _current_spend is not None
- and _current_model_budget_info.max_budget is not None
- and _current_spend > _current_model_budget_info.max_budget
- ):
- raise litellm.BudgetExceededError(
- message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}",
- current_cost=_current_spend,
- max_budget=_current_model_budget_info.max_budget,
- entity_type=Litellm_EntityType.END_USER.value,
- entity_id=end_user_id,
- )
+ max_budget: Final = resolved.budget_config.max_budget
+ if max_budget is None or max_budget < 0:
+ return True
+ current_spend: Final = await self._get_spend_for_model_budget(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ model=model,
+ resolved=resolved,
+ )
+ if current_spend >= max_budget:
+ raise litellm.BudgetExceededError(
+ message=exceeded_message,
+ current_cost=current_spend,
+ max_budget=max_budget,
+ entity_type=entity_type.value,
+ entity_id=entity_id,
+ )
return True
- async def _get_end_user_spend_for_model(
+ async def _get_spend_for_model_budget(
self,
- end_user_id: str,
+ entity_type: Litellm_EntityType,
+ entity_id: str | None,
model: str,
- key_budget_config: BudgetConfig,
- ) -> float | None:
- # 1. model: directly look up `model`
- end_user_model_spend_cache_key = (
- f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}"
- )
- _current_spend = await self.dual_cache.async_get_cache(
- key=end_user_model_spend_cache_key,
- )
+ resolved: ResolvedModelBudget,
+ ) -> float:
+ """Spend charged to this budget in the current window, legacy counter included.
- if _current_spend is None:
- # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
- end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}"
- _current_spend = await self.dual_cache.async_get_cache(
- key=end_user_model_spend_cache_key,
- )
- return _current_spend
-
- async def _get_virtual_key_spend_for_model(
- self,
- user_api_key_hash: str | None,
- model: str,
- key_budget_config: BudgetConfig,
- ) -> float | None:
+ A counter that was never written is zero spend, not unknown spend. The
+ distinction only shows up at a zero-dollar cap, where skipping the
+ comparison would let the strictest possible limit admit every request.
"""
- Get the current spend for a virtual key for a model
-
- Lookup model in this order:
- 1. model: directly look up `model`
- 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
- """
-
- # 1. model: directly look up `model`
- virtual_key_model_spend_cache_key = (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}"
+ spend_key: Final = model_budget_spend_cache_key(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ budget_model=resolved.budget_model,
+ budget_duration=resolved.budget_config.budget_duration,
)
- _current_spend = await self.dual_cache.async_get_cache(
- key=virtual_key_model_spend_cache_key,
+ legacy_spend_key: Final = _legacy_request_model_spend_cache_key(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ model=model,
+ resolved=resolved,
)
+ current_spend: Final = _as_spend(await self._cached_spend(spend_key))
+ if legacy_spend_key is None or legacy_spend_key == spend_key:
+ return current_spend
+ return current_spend + _as_spend(await self._cached_spend(legacy_spend_key))
- if _current_spend is None:
- # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
- # if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview
- virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}"
- _current_spend = await self.dual_cache.async_get_cache(
- key=virtual_key_model_spend_cache_key,
- )
- return _current_spend
-
- def _get_request_model_budget_config(
- self, model: str, internal_model_max_budget: GenericBudgetConfigType
- ) -> BudgetConfig | None:
- """
- Get the budget config for the request model
-
- 1. Check if `model` is in `internal_model_max_budget`
- 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget`
- """
- return internal_model_max_budget.get(model, None) or internal_model_max_budget.get(
- self._get_model_without_custom_llm_provider(model), None
- )
-
- def _get_model_without_custom_llm_provider(self, model: str) -> str:
- if "/" in model:
- return model.split("/")[-1]
- return model
+ async def _cached_spend(self, spend_key: str) -> float | None:
+ return await self.dual_cache.async_get_cache(key=spend_key)
async def async_filter_deployments(
self,
@@ -245,80 +432,63 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
_litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {}
_metadata: Final[dict] = _litellm_params.get("metadata", {}) or {}
- user_api_key_model_max_budget: Final[dict | None] = _metadata.get("user_api_key_model_max_budget", None)
- user_api_key_end_user_model_max_budget: Final[dict | None] = _metadata.get(
- "user_api_key_end_user_model_max_budget", None
- )
- if (user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0) and (
- user_api_key_end_user_model_max_budget is None or len(user_api_key_end_user_model_max_budget) == 0
- ):
- verbose_proxy_logger.debug(
- "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget and user_api_key_end_user_model_max_budget are None or empty."
- )
- return
+ payload_metadata: Final = standard_logging_payload.get("metadata") or {}
- response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
# Use model_group (the user-facing model alias, e.g. "gpt-4o") when
- # available. The enforcement path (is_key_within_model_budget) receives
- # the model name from request_data["model"] which is the model group
- # alias, so the spend tracking cache key must use the same name.
- # Falling back to the deployment-level "model" field preserves
- # behaviour for non-proxy or non-router deployments where model_group
- # is None.
+ # available. The enforcement path receives the model name from
+ # request_data["model"] which is the model group alias, so the spend
+ # tracking cache key must resolve from the same name. Falling back to
+ # the deployment-level "model" field preserves behaviour for non-proxy
+ # or non-router deployments where model_group is None.
model: Final = standard_logging_payload.get("model_group") or standard_logging_payload.get("model")
- virtual_key: Final = standard_logging_payload.get("metadata", {}).get("user_api_key_hash")
- end_user_id = standard_logging_payload.get("end_user") or standard_logging_payload.get("metadata", {}).get(
- "user_api_key_end_user_id"
- )
-
if model is None:
return
- if (
- virtual_key is not None
- and user_api_key_model_max_budget is not None
- and len(user_api_key_model_max_budget) > 0
- ):
- internal_model_max_budget: GenericBudgetConfigType = {}
- for _model, _budget_info in user_api_key_model_max_budget.items():
- internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
- key_budget_config = self._get_request_model_budget_config(
- model=model, internal_model_max_budget=internal_model_max_budget
- )
- if key_budget_config is not None and key_budget_config.budget_duration:
- virtual_spend_key: Final = (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}"
- )
- virtual_start_time_key: Final = f"virtual_key_budget_start_time:{virtual_key}"
- await self._increment_spend_for_key(
- budget_config=key_budget_config,
- spend_key=virtual_spend_key,
- start_time_key=virtual_start_time_key,
- response_cost=response_cost,
- )
+ response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
+ entity_budgets: Final = (
+ (
+ Litellm_EntityType.KEY,
+ payload_metadata.get("user_api_key_hash"),
+ _metadata.get("user_api_key_model_max_budget"),
+ ),
+ (
+ Litellm_EntityType.USER,
+ payload_metadata.get("user_api_key_user_id"),
+ _metadata.get("user_api_key_user_model_max_budget"),
+ ),
+ (
+ Litellm_EntityType.END_USER,
+ standard_logging_payload.get("end_user") or payload_metadata.get("user_api_key_end_user_id"),
+ _metadata.get("user_api_key_end_user_model_max_budget"),
+ ),
+ )
- if (
- end_user_id is not None
- and user_api_key_end_user_model_max_budget is not None
- and len(user_api_key_end_user_model_max_budget) > 0
- ):
- internal_model_max_budget: GenericBudgetConfigType = {}
- for _model, _budget_info in user_api_key_end_user_model_max_budget.items():
- internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
- key_budget_config = self._get_request_model_budget_config(
- model=model, internal_model_max_budget=internal_model_max_budget
+ resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets)
+ if not resolved_budgets:
+ verbose_proxy_logger.debug(
+ "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: "
+ "no key, user or end-user model_max_budget covers model=%s",
+ model,
+ )
+ return
+
+ for entity_type, entity_id, resolved in resolved_budgets:
+ await self._increment_spend_for_key(
+ budget_config=resolved.budget_config,
+ spend_key=model_budget_spend_cache_key(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ budget_model=resolved.budget_model,
+ budget_duration=resolved.budget_config.budget_duration,
+ ),
+ start_time_key=model_budget_start_time_cache_key(
+ entity_type=entity_type,
+ entity_id=entity_id,
+ budget_model=resolved.budget_model,
+ budget_duration=resolved.budget_config.budget_duration,
+ ),
+ response_cost=response_cost,
)
- if key_budget_config is not None and key_budget_config.budget_duration:
- end_user_spend_key: Final = (
- f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}"
- )
- end_user_start_time_key: Final = f"end_user_budget_start_time:{end_user_id}"
- await self._increment_spend_for_key(
- budget_config=key_budget_config,
- spend_key=end_user_spend_key,
- start_time_key=end_user_start_time_key,
- response_cost=response_cost,
- )
if self.dual_cache.redis_cache is not None:
await self._push_in_memory_increments_to_redis()
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 2ec5c34958c..1541b8acfdc 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -1943,6 +1943,8 @@ async def add_litellm_data_to_request(
# Follow same pattern as team and API key budgets
data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend
data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget
+ user_model_budget: Final = user_api_key_dict.user_model_max_budget
+ data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param
data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata)
diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index 9c725c54d08..c2f5b8eeb8b 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
object_permission_cache_key,
user_object_permission_id_cache_key,
)
+from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
from litellm.proxy.management_endpoints.common_daily_activity import (
DailySpendRecord,
@@ -817,6 +818,7 @@ def _build_user_info_response(
keys: list[LiteLLM_VerificationToken] | None,
team_list: list[TeamListResponseObject],
teams_1: list[TeamListResponseObject] | None,
+ model_max_budget_usage: dict[str, dict[str, object]] | None = None,
) -> UserInfoResponse:
"""Create UserInfoResponse while filtering sensitive fields."""
if user_info is None and keys is not None:
@@ -830,6 +832,8 @@ def _build_user_info_response(
if isinstance(_user_info, dict):
_user_info.pop("password", None)
_user_info["metadata"] = _redact_scim_enterprise_metadata(_user_info.get("metadata"))
+ if model_max_budget_usage is not None:
+ _user_info["model_max_budget_usage"] = model_max_budget_usage
return UserInfoResponse(
user_id=user_id,
@@ -864,7 +868,7 @@ async def user_info(
--header 'Authorization: Bearer sk-1234'
```
"""
- from litellm.proxy.proxy_server import prisma_client
+ from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client
try:
user_id = _normalize_user_info_user_id(request=request, user_id=user_id)
@@ -910,6 +914,12 @@ async def user_info(
keys=keys,
team_list=team_list,
teams_1=teams_1,
+ model_max_budget_usage=await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ model_max_budget=getattr(user_info, "model_max_budget", None),
+ cache=model_max_budget_limiter.dual_cache,
+ ),
)
return response_data
@@ -1007,7 +1017,7 @@ async def user_info_v2(
--header 'Authorization: Bearer sk-1234'
```
"""
- from litellm.proxy.proxy_server import prisma_client
+ from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client
try:
if prisma_client is None:
@@ -1062,6 +1072,13 @@ async def user_info_v2(
sso_user_id=user_data.get("sso_user_id"),
teams=user_data.get("teams") or [],
object_permission=user_data.get("object_permission"),
+ model_max_budget=user_data.get("model_max_budget"),
+ model_max_budget_usage=await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_data.get("user_id", user_id),
+ model_max_budget=user_data.get("model_max_budget"),
+ cache=model_max_budget_limiter.dual_cache,
+ ),
)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.user_info_v2(): Exception occured - %s", e)
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 71218d6114b..bf42aeeec05 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -29,6 +29,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
+from litellm.caching.dual_cache import DualCache
from litellm.constants import (
LENGTH_OF_LITELLM_GENERATED_KEY,
LITELLM_PROXY_ADMIN_NAME,
@@ -47,7 +48,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s
rotate_sso_identity_assertions_master_key,
)
from litellm.proxy._types import *
-from litellm.proxy._types import LiteLLM_VerificationToken, hash_token
+from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken, hash_token
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
can_team_access_model,
@@ -73,9 +74,7 @@ from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_k
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
-from litellm.proxy.hooks.model_max_budget_limiter import (
- VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
-)
+from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage
from litellm.proxy.management_endpoints.common_utils import (
_check_passthrough_routes_caller_permission,
_is_user_org_admin_for_team,
@@ -3511,62 +3510,17 @@ async def delete_key_fn(
raise handle_exception_on_proxy(e)
-async def _get_model_max_budget_current_spend(
- api_key_hash: str,
- model: str,
- budget_config: BudgetConfig,
- user_api_key_cache: UserApiKeyCache,
-) -> float:
- virtual_key_model_spend_cache_key = (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}"
- )
- current_spend: float | None = await user_api_key_cache.async_get_cache(
- key=virtual_key_model_spend_cache_key,
- )
- if current_spend is None:
- model_without_prefix: Final = model.split("/")[-1] if "/" in model else model
- virtual_key_model_spend_cache_key = (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:"
- f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}"
- )
- current_spend = await user_api_key_cache.async_get_cache(
- key=virtual_key_model_spend_cache_key,
- )
- try:
- return float(current_spend or 0.0)
- except (TypeError, ValueError):
- return 0.0
-
-
async def _build_model_max_budget_usage(
api_key_hash: str,
model_max_budget: Mapping[str, Mapping[str, object]],
- user_api_key_cache: UserApiKeyCache | None,
+ user_api_key_cache: DualCache | None,
) -> dict[str, dict[str, object]]:
- if user_api_key_cache is None or not model_max_budget:
- return {}
-
- result: Final[dict[str, dict[str, object]]] = {}
- for model, budget_info in model_max_budget.items():
- try:
- budget_config = BudgetConfig.model_validate(budget_info)
- if budget_config.budget_duration is None:
- continue
- duration_in_seconds(budget_config.budget_duration)
- except Exception: # noqa: BLE001
- continue
- spend = await _get_model_max_budget_current_spend(
- api_key_hash=api_key_hash,
- model=model,
- budget_config=budget_config,
- user_api_key_cache=user_api_key_cache,
- )
- result[model] = {
- "current_spend": round(spend, 4),
- "budget_limit": budget_config.max_budget,
- "time_period": budget_config.budget_duration,
- }
- return result
+ return await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id=api_key_hash,
+ model_max_budget=model_max_budget,
+ cache=user_api_key_cache,
+ )
@router.post(
@@ -3596,7 +3550,10 @@ async def info_key_fn_v2(
-d {"keys": ["sk-1", "sk-2", "sk-3"]}
```
"""
- from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
+ from litellm.proxy.proxy_server import (
+ model_max_budget_limiter,
+ prisma_client,
+ )
try:
if prisma_client is None:
@@ -3648,7 +3605,7 @@ async def info_key_fn_v2(
k_dict["model_max_budget_usage"] = await _build_model_max_budget_usage(
api_key_hash=k_token_hash,
model_max_budget=model_max_budget,
- user_api_key_cache=user_api_key_cache,
+ user_api_key_cache=model_max_budget_limiter.dual_cache,
)
filtered_key_info.append(k_dict)
@@ -3707,7 +3664,10 @@ async def info_key_fn(
-H "Authorization: Bearer sk-test-example-key-123"
```
"""
- from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
+ from litellm.proxy.proxy_server import (
+ model_max_budget_limiter,
+ prisma_client,
+ )
try:
if prisma_client is None:
@@ -3760,7 +3720,7 @@ async def info_key_fn(
key_info["model_max_budget_usage"] = await _build_model_max_budget_usage(
api_key_hash=key_token_hash,
model_max_budget=model_max_budget,
- user_api_key_cache=user_api_key_cache,
+ user_api_key_cache=model_max_budget_limiter.dual_cache,
)
# Attach object_permission if object_permission_id is set
@@ -3953,6 +3913,10 @@ async def generate_key_helper_fn(
}
if teams is not None:
user_data["teams"] = teams
+ if model_max_budget:
+ # Only when supplied: the SSO and default-key callers reach this with the
+ # empty default, and writing that would clear an existing user's budgets.
+ user_data["model_max_budget"] = model_max_budget_json
key_data: Final = {
"token": token,
"key_alias": key_alias,
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index d45421489e7..1915a853983 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -64,6 +64,7 @@ from litellm.proxy._types import (
ProxyException,
UserAPIKeyAuth,
)
+from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
@@ -568,6 +569,22 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_metadata["user_api_key"] = user_api_key_dict.api_key
_metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
_metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation
+ # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key
+ # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this
+ # the post-call increment finds nothing and every passthrough request goes untracked and
+ # unenforced. Set after the client merge so a request body cannot supply its own budget.
+ #
+ # Only for the built-in provider routes. `get_model_from_request` returns
+ # None for a user-defined pass-through, deliberately: its body is forwarded
+ # verbatim, so `model` there names an UPSTREAM model rather than a
+ # LiteLLM-managed one. Enforcement is therefore skipped on those routes, and
+ # charging a counter anyway would track spend that nothing can refuse, and
+ # would attribute it to a budget the operator scoped to a LiteLLM model that
+ # merely shares the name.
+ if not request_dispatched_to_pass_through_endpoint(request):
+ _metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget
+ _metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget
+ _metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget
_metadata.update(
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
)
diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
index 55459721906..fcdb3c9246c 100644
--- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
+++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
@@ -2,16 +2,21 @@ import os
import sys
from unittest.mock import AsyncMock, patch
-sys.path.insert(
- 0, os.path.abspath("../..")
-) # Adds the parent directory to the system-path
+sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system-path
import pytest
import litellm
from litellm.caching.caching import DualCache
+from datetime import datetime, timezone
+
+from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.hooks.model_max_budget_limiter import (
+ _budget_model_candidates,
_PROXY_VirtualKeyModelMaxBudgetLimiter,
+ build_model_max_budget_usage,
+ resolve_model_budget,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
@@ -24,41 +29,95 @@ def budget_limiter():
return _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
-# Test _get_model_without_custom_llm_provider
-def test_get_model_without_custom_llm_provider(budget_limiter):
+# Test _budget_model_candidates
+def test_budget_model_candidates():
# Test with custom provider
- assert (
- budget_limiter._get_model_without_custom_llm_provider("openai/gpt-4") == "gpt-4"
- )
+ assert _budget_model_candidates("openai/gpt-4") == ("openai/gpt-4", "gpt-4")
- # Test without custom provider
- assert budget_limiter._get_model_without_custom_llm_provider("gpt-4") == "gpt-4"
+ # Test without custom provider: no duplicate candidate
+ assert _budget_model_candidates("gpt-4") == ("gpt-4",)
-# Test _get_request_model_budget_config
-def test_get_request_model_budget_config(budget_limiter):
- internal_budget = {
- "gpt-4": GenericBudgetInfo(budget_limit=100.0, time_period="1d"),
- "claude-3": GenericBudgetInfo(budget_limit=50.0, time_period="1d"),
+@pytest.mark.parametrize(
+ "model,expected",
+ [
+ (
+ "bedrock/anthropic.claude-opus-4-8",
+ (
+ "bedrock/anthropic.claude-opus-4-8",
+ "anthropic.claude-opus-4-8",
+ "claude-opus-4-8",
+ ),
+ ),
+ (
+ "us.anthropic.claude-opus-4-8",
+ (
+ "us.anthropic.claude-opus-4-8",
+ "anthropic.claude-opus-4-8",
+ "claude-opus-4-8",
+ ),
+ ),
+ (
+ "bedrock/converse/us.amazon.nova-pro-v1:0",
+ (
+ "bedrock/converse/us.amazon.nova-pro-v1:0",
+ "us.amazon.nova-pro-v1:0",
+ "amazon.nova-pro-v1:0",
+ "nova-pro-v1:0",
+ ),
+ ),
+ ],
+)
+def test_budget_model_candidates_reach_the_bedrock_family_name(model, expected):
+ """
+ Bedrock ids carry a dotted vendor segment ("anthropic.", "amazon.") on top of
+ the optional cross-region prefix, so a budget configured under the bare
+ family name would otherwise never match Bedrock traffic: no enforcement and
+ no spend tracking at all.
+ """
+ assert _budget_model_candidates(model) == expected
+
+
+@pytest.mark.parametrize(
+ "model",
+ [
+ "azure/gpt-4.1",
+ "gpt-image-1.5",
+ "not-a-real-model.with.dots",
+ "ft:gpt-4o:acme::abc",
+ ],
+)
+def test_budget_model_candidates_never_split_a_non_bedrock_dotted_name(model):
+ """
+ Most dotted model ids are versions, not Bedrock vendor prefixes. Splitting one
+ would offer a garbage candidate ("gpt-4.1" -> "1") that could collide with an
+ unrelated budget entry, so the split is gated on litellm pricing the model as
+ a Bedrock model.
+ """
+ for candidate in _budget_model_candidates(model):
+ assert candidate in (model, model.split("/")[-1])
+
+
+# Test resolve_model_budget
+def test_resolve_model_budget():
+ model_max_budget = {
+ "gpt-4": {"budget_limit": 100.0, "time_period": "1d"},
+ "claude-3": {"budget_limit": 50.0, "time_period": "1d"},
}
# Test direct model match
- config = budget_limiter._get_request_model_budget_config(
- model="gpt-4", internal_model_max_budget=internal_budget
- )
- assert config.max_budget == 100.0
+ resolved = resolve_model_budget(model="gpt-4", model_max_budget=model_max_budget)
+ assert resolved.budget_model == "gpt-4"
+ assert resolved.budget_config.max_budget == 100.0
- # Test model with provider
- config = budget_limiter._get_request_model_budget_config(
- model="openai/gpt-4", internal_model_max_budget=internal_budget
- )
- assert config.max_budget == 100.0
+ # Test model with provider: the counter is keyed on the CONFIGURED name,
+ # not the request name, so every reader looks it up the same way.
+ resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget)
+ assert resolved.budget_model == "gpt-4"
+ assert resolved.budget_config.max_budget == 100.0
# Test non-existent model
- config = budget_limiter._get_request_model_budget_config(
- model="non-existent", internal_model_max_budget=internal_budget
- )
- assert config is None
+ assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None
# Test is_key_within_model_budget
@@ -72,47 +131,47 @@ async def test_is_key_within_model_budget(budget_limiter):
)
# Test when model is within budget
- with patch.object(
- budget_limiter, "_get_virtual_key_spend_for_model", return_value=50.0
- ):
- assert (
- await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4")
- is True
- )
+ with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0):
+ assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True
# Test when model exceeds budget
- with patch.object(
- budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0
- ):
+ with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0):
with pytest.raises(litellm.BudgetExceededError):
await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4")
# Test model not in budget config
- assert (
- await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent")
- is True
+ assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True
+
+
+# Test _get_spend_for_model_budget
+@pytest.mark.asyncio
+async def test_get_spend_for_model_budget_reads_the_configured_model_key(
+ budget_limiter,
+):
+ from litellm.proxy.hooks.model_max_budget_limiter import (
+ VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
+ model_max_budget = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}
+ # openai/gpt-4 resolves to the configured "gpt-4" entry, so the lookup must
+ # hit the same key async_log_success_event writes.
+ resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget)
-# Test _get_virtual_key_spend_for_model
-@pytest.mark.asyncio
-async def test_get_virtual_key_spend_for_model(budget_limiter):
- budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d")
+ async def _spend(key):
+ return 50.0 if key == f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d" else None
- # Mock cache get
- with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0):
- spend = await budget_limiter._get_virtual_key_spend_for_model(
- user_api_key_hash="test-key", model="gpt-4", key_budget_config=budget_config
- )
- assert spend == 50.0
-
- # Test with provider prefix
- spend = await budget_limiter._get_virtual_key_spend_for_model(
- user_api_key_hash="test-key",
+ with patch.object(budget_limiter.dual_cache, "async_get_cache", side_effect=_spend) as mock_get:
+ spend = await budget_limiter._get_spend_for_model_budget(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id="test-key",
model="openai/gpt-4",
- key_budget_config=budget_config,
+ resolved=resolved,
)
assert spend == 50.0
+ assert [call.kwargs["key"] for call in mock_get.call_args_list] == [
+ f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d",
+ f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:openai/gpt-4:1d",
+ ]
@pytest.mark.asyncio
@@ -138,9 +197,7 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim
"metadata": {"user_api_key_hash": virtual_key},
},
"litellm_params": {
- "metadata": {
- "user_api_key_model_max_budget": user_api_key_model_max_budget
- },
+ "metadata": {"user_api_key_model_max_budget": user_api_key_model_max_budget},
},
}
with patch.object(
@@ -148,15 +205,11 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
- await budget_limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
- assert spend_key == (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}"
- )
+ assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}")
assert call_kwargs["response_cost"] == 0.05
@@ -164,9 +217,7 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim
@pytest.mark.asyncio
async def test_is_end_user_within_model_budget(budget_limiter):
# Test when model is within budget
- with patch.object(
- budget_limiter, "_get_end_user_spend_for_model", return_value=50.0
- ):
+ with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0):
assert (
await budget_limiter.is_end_user_within_model_budget(
"test-user",
@@ -177,9 +228,7 @@ async def test_is_end_user_within_model_budget(budget_limiter):
)
# Test when model exceeds budget
- with patch.object(
- budget_limiter, "_get_end_user_spend_for_model", return_value=150.0
- ):
+ with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0):
with pytest.raises(litellm.BudgetExceededError):
await budget_limiter.is_end_user_within_model_budget(
"test-user",
@@ -198,25 +247,31 @@ async def test_is_end_user_within_model_budget(budget_limiter):
)
-# Test _get_end_user_spend_for_model
+# Test _get_spend_for_model_budget for the end-user scope
@pytest.mark.asyncio
-async def test_get_end_user_spend_for_model(budget_limiter):
- budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d")
+async def test_get_spend_for_end_user_model_budget(budget_limiter):
+ from litellm.proxy.hooks.model_max_budget_limiter import (
+ END_USER_SPEND_CACHE_KEY_PREFIX,
+ )
- # Mock cache get
- with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0):
- spend = await budget_limiter._get_end_user_spend_for_model(
- end_user_id="test-user", model="gpt-4", key_budget_config=budget_config
- )
- assert spend == 50.0
+ model_max_budget = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}
+ resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget)
- # Test with provider prefix
- spend = await budget_limiter._get_end_user_spend_for_model(
- end_user_id="test-user",
+ async def _spend(key):
+ return 50.0 if key == f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:gpt-4:1d" else None
+
+ with patch.object(budget_limiter.dual_cache, "async_get_cache", side_effect=_spend) as mock_get:
+ spend = await budget_limiter._get_spend_for_model_budget(
+ entity_type=Litellm_EntityType.END_USER,
+ entity_id="test-user",
model="openai/gpt-4",
- key_budget_config=budget_config,
+ resolved=resolved,
)
assert spend == 50.0
+ assert [call.kwargs["key"] for call in mock_get.call_args_list] == [
+ f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:gpt-4:1d",
+ f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:openai/gpt-4:1d",
+ ]
@pytest.mark.asyncio
@@ -261,16 +316,12 @@ async def test_async_log_success_event_uses_model_group_for_cache_key(budget_lim
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
- await budget_limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
# The cache key must use the model_group name, NOT the deployment name
- assert spend_key == (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}"
- )
+ assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}")
assert call_kwargs["response_cost"] == 0.10
@@ -310,15 +361,11 @@ async def test_async_log_success_event_falls_back_to_model_when_no_model_group(
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
- await budget_limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
- assert spend_key == (
- f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}"
- )
+ assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}")
@pytest.mark.asyncio
@@ -357,15 +404,11 @@ async def test_async_log_success_event_end_user_uses_model_group(budget_limiter)
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
- await budget_limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
- assert spend_key == (
- f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}"
- )
+ assert spend_key == (f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}")
@pytest.mark.asyncio
@@ -393,9 +436,7 @@ async def test_async_log_success_event_uses_end_user_model_budget_duration(
"metadata": {"user_api_key_end_user_id": end_user_id},
},
"litellm_params": {
- "metadata": {
- "user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget
- },
+ "metadata": {"user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget},
},
}
with patch.object(
@@ -403,15 +444,11 @@ async def test_async_log_success_event_uses_end_user_model_budget_duration(
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
- await budget_limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
- assert spend_key == (
- f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}"
- )
+ assert spend_key == (f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}")
assert call_kwargs["response_cost"] == 0.05
@@ -446,9 +483,7 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config
"_push_in_memory_increments_to_redis",
new_callable=AsyncMock,
) as mock_push:
- await limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_push.assert_awaited_once()
@@ -457,10 +492,7 @@ async def test_get_fallback_model_within_budget_returns_none_without_fallbacks(
budget_limiter,
):
user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={})
- assert (
- await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4")
- is None
- )
+ assert await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") is None
@pytest.mark.asyncio
@@ -472,12 +504,8 @@ async def test_get_fallback_model_within_budget_returns_first_within_budget(
model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}},
budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]},
)
- with patch.object(
- budget_limiter, "_get_virtual_key_spend_for_model", return_value=1.0
- ):
- result = await budget_limiter.get_fallback_model_within_budget(
- user_api_key, "gpt-4"
- )
+ with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=1.0):
+ result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4")
assert result == "gpt-4o-mini"
@@ -494,17 +522,15 @@ async def test_get_fallback_model_within_budget_skips_exhausted_fallback(
budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]},
)
- async def _spend_for_model(user_api_key_hash, model, key_budget_config):
- return 150.0 if model == "gpt-4o-mini" else 1.0
+ async def _spend_for_model(entity_type, entity_id, model, resolved):
+ return 150.0 if resolved.budget_model == "gpt-4o-mini" else 1.0
with patch.object(
budget_limiter,
- "_get_virtual_key_spend_for_model",
+ "_get_spend_for_model_budget",
side_effect=_spend_for_model,
):
- result = await budget_limiter.get_fallback_model_within_budget(
- user_api_key, "gpt-4"
- )
+ result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4")
assert result == "claude-haiku"
@@ -520,12 +546,8 @@ async def test_get_fallback_model_within_budget_returns_none_when_chain_exhauste
},
budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]},
)
- with patch.object(
- budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0
- ):
- result = await budget_limiter.get_fallback_model_within_budget(
- user_api_key, "gpt-4"
- )
+ with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0):
+ result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4")
assert result is None
@@ -554,7 +576,761 @@ async def test_async_log_success_event_skips_redis_push_without_redis(budget_lim
"_push_in_memory_increments_to_redis",
new_callable=AsyncMock,
) as mock_push:
- await budget_limiter.async_log_success_event(
- kwargs, response_obj=None, start_time=None, end_time=None
- )
+ await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
mock_push.assert_not_awaited()
+
+
+def _success_kwargs(
+ *,
+ model_group,
+ deployment_model=None,
+ response_cost=0.5,
+ key_hash=None,
+ key_model_max_budget=None,
+ user_id=None,
+ user_model_max_budget=None,
+ end_user_id=None,
+ end_user_model_max_budget=None,
+):
+ return {
+ "standard_logging_object": {
+ "response_cost": response_cost,
+ "model": deployment_model or model_group,
+ "model_group": model_group,
+ "end_user": end_user_id,
+ "metadata": {
+ "user_api_key_hash": key_hash,
+ "user_api_key_user_id": user_id,
+ "user_api_key_end_user_id": end_user_id,
+ },
+ },
+ "litellm_params": {
+ "metadata": {
+ "user_api_key_model_max_budget": key_model_max_budget,
+ "user_api_key_user_model_max_budget": user_model_max_budget,
+ "user_api_key_end_user_model_max_budget": end_user_model_max_budget,
+ },
+ },
+ }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "request_model",
+ ["gpt-4", "openai/gpt-4"],
+ ids=["request_model_matches_budget_key", "request_model_carries_provider_prefix"],
+)
+async def test_logged_spend_is_visible_to_key_info_usage_and_enforcement(request_model):
+ """
+ The counter written post-call, the counter enforcement reads and the counter
+ /key/info reports must be one and the same, including when the request model
+ is not byte-identical to the configured budget key.
+
+ Regression: the increment used to be keyed on the REQUEST model while
+ /key/info only ever looked up the CONFIGURED model, so a key could be
+ actively blocked at 429 while reporting current_spend 0.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ key_hash = "vk-hash"
+ model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group=request_model,
+ response_cost=0.75,
+ key_hash=key_hash,
+ key_model_max_budget=model_max_budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ usage = await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id=key_hash,
+ model_max_budget=model_max_budget,
+ cache=dual_cache,
+ )
+ assert usage == {
+ "gpt-4": {
+ "current_spend": 0.75,
+ "budget_limit": 1.0,
+ "time_period": "1d",
+ }
+ }
+
+ user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget)
+ # Still under the 1.0 limit.
+ assert await limiter.is_key_within_model_budget(user_api_key, request_model) is True
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group=request_model,
+ response_cost=0.75,
+ key_hash=key_hash,
+ key_model_max_budget=model_max_budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_key_within_model_budget(user_api_key, request_model)
+
+ usage_after = await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id=key_hash,
+ model_max_budget=model_max_budget,
+ cache=dual_cache,
+ )
+ assert usage_after["gpt-4"]["current_spend"] == 1.5
+
+
+@pytest.mark.asyncio
+async def test_user_model_budget_is_tracked_and_enforced():
+ """
+ An internal user's own model_max_budget must be incremented post-call and
+ enforced, independently of any key-level budget.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ user_id = "user-1"
+ user_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}
+
+ assert (
+ await limiter.is_user_within_model_budget(
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ model="openai/gpt-4",
+ )
+ is True
+ )
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group="openai/gpt-4",
+ response_cost=1.5,
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ assert await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ model_max_budget=user_model_max_budget,
+ cache=dual_cache,
+ ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}}
+
+ with pytest.raises(litellm.BudgetExceededError) as exc:
+ await limiter.is_user_within_model_budget(
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ model="openai/gpt-4",
+ )
+ assert exc.value.entity_type == Litellm_EntityType.USER.value
+
+
+@pytest.mark.asyncio
+async def test_user_model_budget_counter_is_separate_from_the_key_counter():
+ """
+ A key budget and a user budget over the same model are two independent
+ counters, so one request must charge each exactly once.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group="gpt-4",
+ response_cost=2.0,
+ key_hash="vk-hash",
+ key_model_max_budget=model_max_budget,
+ user_id="user-1",
+ user_model_max_budget=model_max_budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-hash:gpt-4:1d") == 2.0
+ assert await dual_cache.async_get_cache(key="user_model_spend:user-1:gpt-4:1d") == 2.0
+
+
+@pytest.mark.asyncio
+async def test_two_models_on_one_key_do_not_share_a_budget_window():
+ """
+ A key budgeting two models over different periods must own one window start
+ per model: a shared start lets the shorter period restart the longer one.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ model_max_budget = {
+ "gpt-4": {"budget_limit": 10.0, "time_period": "1d"},
+ "claude-3": {"budget_limit": 10.0, "time_period": "30d"},
+ }
+
+ start_time_keys = []
+ with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment:
+ for model in ("gpt-4", "claude-3"):
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group=model,
+ key_hash="vk-hash",
+ key_model_max_budget=model_max_budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+ start_time_keys = [call.kwargs["start_time_key"] for call in mock_increment.call_args_list]
+
+ assert start_time_keys == [
+ "virtual_key_budget_start_time:vk-hash:gpt-4:1d",
+ "virtual_key_budget_start_time:vk-hash:claude-3:30d",
+ ]
+ assert len(set(start_time_keys)) == 2
+
+
+@pytest.mark.asyncio
+async def test_no_increment_when_no_scope_budgets_the_model():
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment:
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group="gpt-4",
+ key_hash="vk-hash",
+ key_model_max_budget={"claude-3": {"budget_limit": 1.0, "time_period": "1d"}},
+ user_id="user-1",
+ user_model_max_budget={},
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+ mock_increment.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_build_model_max_budget_usage_skips_unusable_entries():
+ """A malformed or period-less entry must be omitted, not crash the report."""
+ dual_cache = DualCache()
+ await dual_cache.async_set_cache(key="virtual_key_spend:vk:gpt-4:1d", value=3.0)
+
+ usage = await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id="vk",
+ model_max_budget={
+ "gpt-4": {"budget_limit": 10.0, "time_period": "1d"},
+ "no-period": {"budget_limit": 10.0},
+ "bad-period": {"budget_limit": 10.0, "time_period": "not-a-duration"},
+ },
+ cache=dual_cache,
+ )
+ assert usage == {"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d"}}
+
+
+@pytest.mark.asyncio
+async def test_bedrock_traffic_charges_the_bare_family_name_budget():
+ """
+ The reported case: a budget configured as "claude-opus-4-8" with traffic on
+ "bedrock/anthropic.claude-opus-4-8". Before the fix nothing matched, so spend
+ was never tracked and the budget was never enforced no matter how far over it
+ the key went.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ key_hash = "vk-hash"
+ model_max_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}
+ user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget)
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group="bedrock/anthropic.claude-opus-4-8",
+ response_cost=1.5,
+ key_hash=key_hash,
+ key_model_max_budget=model_max_budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ assert await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id=key_hash,
+ model_max_budget=model_max_budget,
+ cache=dual_cache,
+ ) == {
+ "claude-opus-4-8": {
+ "current_spend": 1.5,
+ "budget_limit": 1.0,
+ "time_period": "18h",
+ }
+ }
+
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_key_within_model_budget(user_api_key, "bedrock/anthropic.claude-opus-4-8")
+
+
+@pytest.mark.asyncio
+async def test_user_model_budget_window_resets_when_the_period_elapses():
+ """
+ A monthly user budget must start a fresh window once the period elapses,
+ and the window start must be scoped to that one budget model so a second
+ model on a shorter period cannot drag it forward.
+ """
+ from litellm.proxy.hooks.model_max_budget_limiter import (
+ model_budget_spend_cache_key,
+ model_budget_start_time_cache_key,
+ )
+
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ user_id = "user-1"
+ user_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}
+ spend_key = model_budget_spend_cache_key(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ budget_model="gpt-4",
+ budget_duration="1mo",
+ )
+ start_time_key = model_budget_start_time_cache_key(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ budget_model="gpt-4",
+ budget_duration="1mo",
+ )
+
+ kwargs = _success_kwargs(
+ model_group="gpt-4",
+ response_cost=1.5,
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ )
+ await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
+ assert await dual_cache.async_get_cache(key=spend_key) == 1.5
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_user_within_model_budget(
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ model="gpt-4",
+ )
+
+ # Age the window past its period. The next charge opens a new window rather
+ # than adding to the exhausted one.
+ elapsed = duration_in_seconds("1mo") + 60
+ await dual_cache.async_set_cache(
+ key=start_time_key,
+ value=datetime.now(timezone.utc).timestamp() - elapsed,
+ ttl=elapsed,
+ )
+
+ await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None)
+ assert await dual_cache.async_get_cache(key=spend_key) == 1.5
+ assert await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ model_max_budget=user_model_max_budget,
+ cache=dual_cache,
+ ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}}
+
+
+@pytest.mark.asyncio
+async def test_a_zero_dollar_cap_blocks_the_model():
+ """
+ 0 is the operator saying "nobody may spend anything on this model", which is
+ the strictest cap expressible, not the absence of one. Skipping it on
+ falsiness turned the strictest setting into no setting at all, so the model
+ stayed wide open. The dashboard editor can produce this value, so it has to
+ mean something.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ key = UserAPIKeyAuth(
+ token="hash-zero",
+ model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}},
+ )
+
+ with pytest.raises(litellm.BudgetExceededError) as exc:
+ await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4")
+ assert exc.value.max_budget == 0
+
+
+@pytest.mark.asyncio
+async def test_a_zero_dollar_cap_is_reported_as_a_cap_not_as_absent():
+ """The usage endpoints must show the 0 too, or an operator cannot see the block they configured."""
+ assert await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id="hash-zero",
+ model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}},
+ cache=DualCache(),
+ ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d"}}
+
+
+@pytest.mark.asyncio
+async def test_spend_exactly_at_the_cap_is_refused():
+ """
+ Spending the whole budget exhausts it. `>` let a caller sit exactly on the
+ limit and keep going, and every sibling budget check in the codebase
+ (RouterBudgetLimiting, the key and team budget checks) uses `>=`.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ budget = {"gpt-4": {"budget_limit": 2.0, "time_period": "1d"}}
+ key = UserAPIKeyAuth(token="hash-exact", model_max_budget=budget)
+
+ await limiter.async_log_success_event(
+ _success_kwargs(model_group="gpt-4", response_cost=2.0, key_hash="hash-exact", key_model_max_budget=budget),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4")
+
+
+@pytest.mark.asyncio
+async def test_usage_report_reads_every_counter_in_one_batched_lookup():
+ """
+ model_max_budget is caller-supplied and unbounded in size, so one cache
+ coroutine per configured model let a large map fan out into an unbounded
+ number of concurrent lookups on an endpoint anyone holding the key can call.
+ One batched read keeps it to a single round trip whatever the map's size.
+ """
+ dual_cache = DualCache()
+ budget = {f"model-{i}": {"budget_limit": 1.0, "time_period": "1d"} for i in range(50)}
+
+ with (
+ patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 50)) as batched,
+ patch.object(dual_cache, "async_get_cache", new=AsyncMock()) as single,
+ ):
+ usage = await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id="hash-many",
+ model_max_budget=budget,
+ cache=dual_cache,
+ )
+
+ assert batched.await_count == 1
+ assert len(batched.await_args.kwargs["keys"]) == 50
+ assert single.await_count == 0
+ assert len(usage) == 50
+
+
+@pytest.mark.asyncio
+async def test_usage_report_survives_a_batch_lookup_that_returns_nothing():
+ """
+ async_batch_get_cache swallows its own failures and returns None. Zipping
+ that against the budgets would raise and take the whole /key/info response
+ with it, so an unusable result has to read as a miss instead.
+ """
+ dual_cache = DualCache()
+ with patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=None)):
+ assert await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id="hash-none",
+ model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}},
+ cache=dual_cache,
+ ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d"}}
+
+
+@pytest.mark.asyncio
+async def test_one_malformed_scope_does_not_abort_the_other_scopes():
+ """
+ Every scope is resolved before any of them is incremented, so a single
+ unusable entry used to raise out of resolution and leave the key counter
+ unwritten too. The key's budget is well formed here and must still be
+ charged despite the user's entry being garbage.
+ """
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ key_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group="gpt-4",
+ response_cost=0.25,
+ key_hash="hash-mixed",
+ key_model_max_budget=key_budget,
+ user_id="user-mixed",
+ user_model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}},
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ assert await build_model_max_budget_usage(
+ entity_type=Litellm_EntityType.KEY,
+ entity_id="hash-mixed",
+ model_max_budget=key_budget,
+ cache=dual_cache,
+ ) == {"gpt-4": {"current_spend": 0.25, "budget_limit": 10.0, "time_period": "1d"}}
+
+
+@pytest.mark.asyncio
+async def test_an_unusable_budget_entry_is_not_enforced_instead_of_raising():
+ """
+ A config typo must not turn every request for that model into a 500. It
+ cannot be keyed, so it cannot be enforced; the write path rejects these, so
+ reaching here means config.yaml or a direct DB edit.
+ """
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+ key = UserAPIKeyAuth(
+ token="hash-malformed",
+ model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}},
+ )
+
+ assert await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") is True
+
+
+def test_resolve_model_budget_returns_none_for_an_unusable_entry():
+ assert (
+ resolve_model_budget(
+ model="gpt-4",
+ model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}},
+ )
+ is None
+ )
+
+
+def test_a_malformed_specific_entry_does_not_hide_a_usable_family_budget():
+ """
+ The candidate chain is most-specific-first and already falls through an entry
+ that is ABSENT. An entry that will not parse is indistinguishable from absent
+ as far as enforcement goes, so it has to fall through too: otherwise one bad
+ provider-prefixed entry silently disables the valid bare-family budget sitting
+ next to it, and the model goes uncapped.
+ """
+ resolved = resolve_model_budget(
+ model="openai/gpt-4",
+ model_max_budget={
+ "openai/gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"},
+ "gpt-4": {"budget_limit": 7.0, "time_period": "1d"},
+ },
+ )
+
+ assert resolved is not None
+ assert resolved.budget_model == "gpt-4"
+ assert resolved.budget_config.max_budget == 7.0
+
+
+@pytest.mark.asyncio
+async def test_a_malformed_specific_entry_still_enforces_the_family_budget():
+ """The fall-through has to reach enforcement, not just resolution."""
+ dual_cache = DualCache()
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+ budget = {
+ "openai/gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"},
+ "gpt-4": {"budget_limit": 1.0, "time_period": "1d"},
+ }
+ key = UserAPIKeyAuth(token="hash-fallthrough", model_max_budget=budget)
+
+ await limiter.async_log_success_event(
+ _success_kwargs(
+ model_group="openai/gpt-4",
+ response_cost=2.0,
+ key_hash="hash-fallthrough",
+ key_model_max_budget=budget,
+ ),
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_key_within_model_budget(user_api_key_dict=key, model="openai/gpt-4")
+
+
+def test_documented_budget_spelling_survives_model_validate():
+ """
+ `budget_limit` / `time_period` are the spelling the docs, the CRUD endpoints
+ and the dashboard editor all use, and BudgetConfig maps them onto
+ `max_budget` / `budget_duration` inside its `__init__`.
+
+ Pydantic v2 normally bypasses a custom `__init__` in `model_validate`, and
+ this code path validates rather than constructing. It works today, but that
+ is a property of the installed Pydantic rather than of anything in this
+ repository, so an upgrade could silently stop applying the mapping and
+ quietly disable every budget written in the documented spelling. Pinned here
+ so that becomes a red test instead of an outage.
+ """
+ from litellm.types.utils import BudgetConfig
+
+ validated = BudgetConfig.model_validate({"budget_limit": 5, "time_period": "1d"})
+ assert validated.max_budget == 5.0
+ assert validated.budget_duration == "1d"
+
+ # Control: an unrecognised key must NOT populate max_budget, or the assertion
+ # above would also pass against a model that accepted anything at all.
+ ignored = BudgetConfig.model_validate({"bogus_limit": 5, "time_period": "1d"})
+ assert ignored.max_budget is None
+
+
+def test_resolution_accepts_both_documented_spellings():
+ """The resolver is what enforcement, tracking and reporting all go through."""
+ for budget in (
+ {"gpt-4": {"budget_limit": 5, "time_period": "1d"}},
+ {"gpt-4": {"max_budget": 5, "budget_duration": "1d"}},
+ ):
+ resolved = resolve_model_budget(model="gpt-4", model_max_budget=budget)
+ assert resolved is not None, f"{budget} resolved to nothing"
+ assert resolved.budget_config.max_budget == 5.0
+ assert resolved.budget_config.budget_duration == "1d"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "entity_type, prefix",
+ [
+ (Litellm_EntityType.KEY, "virtual_key_spend"),
+ (Litellm_EntityType.END_USER, "end_user_model_spend"),
+ ],
+)
+async def test_a_pre_upgrade_counter_keyed_on_the_request_model_still_enforces(entity_type, prefix):
+ """An upgrading proxy must not hand out a second allowance for the window it is already in.
+
+ Before the counter key moved to the configured budget model, spend for a
+ request on `openai/gpt-4` against a budget configured as `gpt-4` was both
+ written to and enforced on `{prefix}:{id}:openai/gpt-4:1d`. Reading only the
+ configured-model key finds that counter empty and admits another full budget
+ until the window expires.
+ """
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+ model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
+ await limiter.dual_cache.async_set_cache(key=f"{prefix}:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400)
+
+ with pytest.raises(litellm.BudgetExceededError) as exc_info:
+ if entity_type == Litellm_EntityType.KEY:
+ await limiter.is_key_within_model_budget(
+ user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget),
+ model="openai/gpt-4",
+ )
+ else:
+ await limiter.is_end_user_within_model_budget(
+ end_user_id="entity-1",
+ end_user_model_max_budget=model_max_budget,
+ model="openai/gpt-4",
+ )
+ assert exc_info.value.current_cost == 25.0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "legacy_spend, current_spend, expect_blocked",
+ [(6.0, 5.0, True), (2.0, 3.0, False)],
+)
+async def test_the_pre_upgrade_and_post_upgrade_counters_add_up_over_one_window(
+ legacy_spend, current_spend, expect_blocked
+):
+ """The two counters hold disjoint halves of one window, so the window's spend is their sum.
+
+ Nothing writes the request-model spelling once this version is running, so
+ the legacy counter is frozen at whatever the previous version charged and
+ the configured-model counter carries everything since. Either one alone
+ under-reports the window: 6 + 5 is over a cap of 10 that neither half
+ reaches on its own.
+ """
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+ model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
+ await limiter.dual_cache.async_set_cache(
+ key="virtual_key_spend:entity-1:openai/gpt-4:1d", value=legacy_spend, ttl=86400
+ )
+ await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:gpt-4:1d", value=current_spend, ttl=86400)
+
+ async def enforce():
+ return await limiter.is_key_within_model_budget(
+ user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget),
+ model="openai/gpt-4",
+ )
+
+ if expect_blocked:
+ with pytest.raises(litellm.BudgetExceededError) as exc_info:
+ await enforce()
+ assert exc_info.value.current_cost == legacy_spend + current_spend
+ else:
+ assert await enforce() is True
+
+
+@pytest.mark.asyncio
+async def test_the_configured_model_counter_is_never_counted_twice():
+ """When the request names the budget exactly there is no legacy counter, only the one key.
+
+ Both keys are `virtual_key_spend:entity-1:gpt-4:1d` here, so a lookup that
+ added them without noticing would charge 12 against a cap of 10 and refuse a
+ key that has spent 6.
+ """
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+ await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:gpt-4:1d", value=6.0, ttl=86400)
+
+ assert (
+ await limiter.is_key_within_model_budget(
+ user_api_key_dict=UserAPIKeyAuth(
+ token="entity-1",
+ model_max_budget={"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}},
+ ),
+ model="gpt-4",
+ )
+ is True
+ )
+
+
+@pytest.mark.asyncio
+async def test_the_pre_upgrade_counter_is_no_longer_read_a_window_after_start_up(monkeypatch):
+ """The carry is bounded, so it cannot become a permanent second lookup on every request.
+
+ A counter written by the previous version belongs to a window that was
+ already open when this process replaced it, so once a full window has passed
+ since start-up there is nothing left for the lookup to find.
+ """
+ import litellm.proxy.hooks.model_max_budget_limiter as limiter_module
+
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+ model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
+ await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400)
+ user_api_key = UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget)
+
+ # Control: within the first window since start-up the same counter blocks,
+ # so the assertion below cannot pass against a lookup that never worked.
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_key_within_model_budget(user_api_key_dict=user_api_key, model="openai/gpt-4")
+
+ monkeypatch.setattr(limiter_module, "_PROCESS_STARTED_AT", limiter_module.time.monotonic() - 86401)
+ assert await limiter.is_key_within_model_budget(user_api_key_dict=user_api_key, model="openai/gpt-4") is True
+
+
+@pytest.mark.asyncio
+async def test_the_user_scope_has_no_pre_upgrade_counter_to_carry():
+ """The user scope is introduced by this change, so a request-model key under it is not one of ours.
+
+ Reading one would invent a counter no previous version ever wrote, which is
+ the opposite of preserving one.
+ """
+ limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+ model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
+ await limiter.dual_cache.async_set_cache(key="user_model_spend:u1:openai/gpt-4:1d", value=25.0, ttl=86400)
+
+ assert (
+ await limiter.is_user_within_model_budget(
+ user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4"
+ )
+ is True
+ )
+
+ # Control: the same overspend under the key this scope does own must block,
+ # or the assertion above would pass against a scope that enforces nothing.
+ await limiter.dual_cache.async_set_cache(key="user_model_spend:u1:gpt-4:1d", value=25.0, ttl=86400)
+ with pytest.raises(litellm.BudgetExceededError):
+ await limiter.is_user_within_model_budget(
+ user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4"
+ )
diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py
index 58dbe3ad370..e9566254dbc 100644
--- a/tests/proxy_unit_tests/test_user_api_key_auth.py
+++ b/tests/proxy_unit_tests/test_user_api_key_auth.py
@@ -7,9 +7,7 @@ import sys
import litellm.proxy
import litellm.proxy.proxy_server
-sys.path.insert(
- 0, os.path.abspath("../..")
-) # Adds the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path
from typing import Dict, List, Optional
from unittest.mock import MagicMock, patch, AsyncMock
@@ -50,9 +48,7 @@ class Request:
), # Request with no client IP should not be allowed
],
)
-def test_check_valid_ip(
- allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool
-):
+def test_check_valid_ip(allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool):
from litellm.proxy.auth.auth_utils import _check_valid_ip
request = Request(client_ip)
@@ -121,9 +117,7 @@ async def test_check_blocked_team():
last_refreshed_at=time.time(),
)
await asyncio.sleep(1)
- team_obj = LiteLLM_TeamTableCachedObj(
- team_id=_team_id, blocked=False, last_refreshed_at=time.time()
- )
+ team_obj = LiteLLM_TeamTableCachedObj(team_id=_team_id, blocked=False, last_refreshed_at=time.time())
hashed_token = hash_token(user_key)
print(f"STORING TOKEN UNDER KEY={hashed_token}")
user_api_key_cache.set_cache(key=hashed_token, value=valid_token)
@@ -173,9 +167,7 @@ async def test_team_object_has_object_permission_id():
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
- with patch(
- "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
- ) as mock_common_checks:
+ with patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock) as mock_common_checks:
mock_common_checks.return_value = True
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
@@ -200,9 +192,7 @@ async def test_returned_user_api_key_auth(user_role, expected_role):
from datetime import datetime
new_obj = await _return_user_api_key_auth_obj(
- user_obj=LiteLLM_UserTable(
- user_role=user_role, user_id="", max_budget=None, user_email=""
- ),
+ user_obj=LiteLLM_UserTable(user_role=user_role, user_id="", max_budget=None, user_email=""),
api_key="hello-world",
parent_otel_span=None,
valid_token_dict={},
@@ -258,9 +248,7 @@ async def test_aaauser_personal_budgets(key_ownership):
spend=20,
)
- user_obj = LiteLLM_UserTable(
- user_id=_user_id, spend=11, max_budget=10, user_email=""
- )
+ user_obj = LiteLLM_UserTable(user_id=_user_id, spend=11, max_budget=10, user_email="")
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
user_api_key_cache.set_cache(key="{}".format(_user_id), value=user_obj)
@@ -273,10 +261,7 @@ async def test_aaauser_personal_budgets(key_ownership):
test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache")
- assert (
- test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth)
- == valid_token
- )
+ assert test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth) == valid_token
if key_ownership == "user_key":
with pytest.raises(ProxyException) as exc_info:
@@ -311,9 +296,7 @@ async def test_user_api_key_auth_fails_with_prohibited_params(prohibited_param):
request.body = return_body
try:
- response = await user_api_key_auth(
- request=request, api_key="Bearer " + user_key
- )
+ response = await user_api_key_auth(request=request, api_key="Bearer " + user_key)
except Exception as e:
print("error str=", str(e))
error_message = str(e.message)
@@ -519,9 +502,7 @@ def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api
verbose_proxy_logger.setLevel(logging.DEBUG)
request = MagicMock(spec=Request)
request.headers = headers
- api_key = get_api_key_from_custom_header(
- request=request, custom_litellm_key_header_name=custom_header_name
- )
+ api_key = get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=custom_header_name)
assert api_key == expected_api_key
@@ -572,9 +553,7 @@ from litellm.proxy._types import LitellmUserRoles
(LitellmUserRoles.TEAM, "1234", "1234", True),
],
)
-def test_allowed_route_inside_route(
- user_role, auth_user_id, requested_user_id, expected_result
-):
+def test_allowed_route_inside_route(user_role, auth_user_id, requested_user_id, expected_result):
from litellm.proxy.auth.auth_checks import allowed_route_check_inside_route
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
@@ -715,9 +694,7 @@ async def test_soft_budget_alert():
try:
# Call user_api_key_auth
- response = await user_api_key_auth(
- request=request, api_key="Bearer " + user_key
- )
+ response = await user_api_key_auth(request=request, api_key="Bearer " + user_key)
# Assert the request was allowed (no exception raised)
assert response is not None
@@ -883,9 +860,7 @@ async def test_user_api_key_auth_websocket():
mock_websocket.url = URL(url="/ws")
# Mock the return value of `user_api_key_auth` when it's called within the `user_api_key_auth_websocket` function
- with patch(
- "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True
- ) as mock_user_api_key_auth:
+ with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth:
# Make the call to the WebSocket function
await user_api_key_auth_websocket(mock_websocket)
@@ -896,17 +871,11 @@ async def test_user_api_key_auth_websocket():
request_arg = mock_user_api_key_auth.call_args.kwargs["request"]
# Verify that the request has headers set
- assert hasattr(
- request_arg, "headers"
- ), "Request object should have headers attribute"
- assert (
- "authorization" in request_arg.headers
- ), "Request headers should contain authorization"
+ assert hasattr(request_arg, "headers"), "Request object should have headers attribute"
+ assert "authorization" in request_arg.headers, "Request headers should contain authorization"
assert request_arg.headers["authorization"] == "Bearer some_api_key"
- assert (
- mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key"
- )
+ assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key"
@pytest.mark.asyncio
@@ -929,9 +898,7 @@ async def test_user_api_key_auth_websocket_carries_asgi_path():
}
mock_websocket.url = URL(url="/v1/realtime")
- with patch(
- "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True
- ) as mock_user_api_key_auth:
+ with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth:
await user_api_key_auth_websocket(mock_websocket)
request_arg = mock_user_api_key_auth.call_args.kwargs["request"]
@@ -1127,9 +1094,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
)
request._url = URL(url="/team/new")
- monkeypatch.setattr(
- litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True}
- )
+ monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True})
# Initialize jwt_handler with a default LiteLLM_JWTAuth so that the
# virtual_key_claim_field check in user_api_key_auth doesn't fail with
@@ -1158,9 +1123,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
):
try:
await user_api_key_auth(request=request, api_key="Bearer fake.jwt.token")
- pytest.fail(
- "Expected this call to fail. Non-admin user should not access team routes."
- )
+ pytest.fail("Expected this call to fail. Non-admin user should not access team routes.")
except ProxyException as e:
print("e", e)
assert "Only proxy admin can be used to generate" in str(e.message)
@@ -1220,9 +1183,7 @@ async def test_user_api_key_from_query_param():
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
user_key = "sk-query-1234"
- user_api_key_cache.set_cache(
- key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key))
- )
+ user_api_key_cache.set_cache(key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key)))
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@@ -1235,9 +1196,7 @@ async def test_user_api_key_from_query_param():
"query_string": f"alt=sse&key={user_key}".encode(),
}
)
- request._url = URL(
- url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}"
- )
+ request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}")
async def return_body():
return b"{}"
@@ -1246,3 +1205,592 @@ async def test_user_api_key_from_query_param():
valid_token = await user_api_key_auth(request=request, api_key="")
assert valid_token.token == hash_token(user_key)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "user_id,user_model_max_budget,expected_calls",
+ [
+ ("u-1", {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, 1),
+ ("u-1", {}, 0),
+ ("u-1", None, 0),
+ (None, {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, 0),
+ ],
+ ids=["enforced", "empty_budget", "no_budget", "no_user_id"],
+)
+async def test_check_user_model_budget(user_id, user_model_max_budget, expected_calls):
+ """
+ An internal user's model_max_budget must reach the limiter. Before this it was
+ stored on LiteLLM_UserTable, accepted by /user/new and /user/update, and read
+ by nothing, so a user-level per-model budget never blocked anything.
+ """
+ from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget
+
+ calls = []
+
+ class _Limiter:
+ async def is_user_within_model_budget(self, user_id, user_model_max_budget, model):
+ calls.append((user_id, user_model_max_budget, model))
+ return True
+
+ valid_token = UserAPIKeyAuth(
+ token="hash",
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ )
+ await _check_user_model_budget(
+ valid_token=valid_token,
+ model_max_budget_limiter=_Limiter(),
+ models=["gpt-4"],
+ )
+ assert len(calls) == expected_calls
+ if expected_calls:
+ assert calls[0] == ("u-1", user_model_max_budget, "gpt-4")
+
+
+@pytest.mark.asyncio
+async def test_user_model_max_budget_is_threaded_onto_the_auth_object():
+ """
+ The limiter can only enforce what auth carries. Regression for the user row's
+ model_max_budget being dropped on the way into UserAPIKeyAuth.
+ """
+ from datetime import datetime
+
+ from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj
+
+ budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}
+ user_obj = LiteLLM_UserTable(
+ user_id="u-1",
+ max_budget=None,
+ spend=0.0,
+ user_email=None,
+ models=[],
+ model_max_budget=budget,
+ )
+
+ auth_obj = await _return_user_api_key_auth_obj(
+ user_obj=user_obj,
+ api_key="sk-1234",
+ parent_otel_span=None,
+ valid_token_dict={"token": "hash"},
+ route="/chat/completions",
+ start_time=datetime.now(),
+ user_role=LitellmUserRoles.INTERNAL_USER,
+ )
+ assert auth_obj.user_model_max_budget == budget
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "over_budget,expect_refusal",
+ [(True, True), (False, False)],
+ ids=["over_budget_is_refused", "under_budget_is_served"],
+)
+async def test_user_model_budget_is_enforced_through_user_api_key_auth(over_budget, expect_refusal):
+ """
+ Drive the real auth entry point, not the helper.
+
+ The user's model_max_budget lives on the user row, and the joint
+ verification-token view auth builds its token from does not carry it. A test
+ that only exercises the helper passes while the whole path is inert, so this
+ one goes through user_api_key_auth with a key that has no per-model budget of
+ its own and asserts the USER's budget decides the outcome.
+ """
+ from fastapi import Request
+ from starlette.datastructures import URL
+
+ from litellm.proxy._types import LiteLLM_UserTable, Litellm_EntityType, UserAPIKeyAuth
+ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+ from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key
+ from litellm.proxy.proxy_server import (
+ hash_token,
+ model_max_budget_limiter,
+ user_api_key_cache,
+ )
+
+ user_id = "user-model-budget"
+ model = "gpt-4o"
+ key = "sk-user-model-budget"
+ hashed = hash_token(key)
+ user_model_max_budget = {model: {"budget_limit": 1.0, "time_period": "1mo"}}
+
+ setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
+ setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+ setattr(litellm.proxy.proxy_server, "prisma_client", "present")
+
+ await user_api_key_cache.async_set_cache(
+ key=hashed,
+ value=UserAPIKeyAuth(token=hashed, user_id=user_id, models=[], model_max_budget={}),
+ model_type=UserAPIKeyAuth,
+ )
+ await model_max_budget_limiter.dual_cache.async_set_cache(
+ key=model_budget_spend_cache_key(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ budget_model=model,
+ budget_duration="1mo",
+ ),
+ value=5.0 if over_budget else 0.25,
+ ttl=600,
+ )
+
+ request = Request(scope={"type": "http"})
+ request._url = URL(url="/chat/completions")
+
+ async def return_body():
+ return f'{{"model": "{model}"}}'.encode()
+
+ request.body = return_body
+
+ async def fake_get_user_object(**kwargs):
+ return LiteLLM_UserTable(
+ user_id=user_id,
+ max_budget=None,
+ spend=0.0,
+ user_email=None,
+ models=[],
+ model_max_budget=user_model_max_budget,
+ )
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.get_user_object",
+ new=fake_get_user_object,
+ ):
+ if expect_refusal:
+ with pytest.raises(Exception) as exc:
+ await user_api_key_auth(request=request, api_key="Bearer " + key)
+ assert "budget" in str(exc.value).lower()
+ assert user_id in str(exc.value)
+ else:
+ result = await user_api_key_auth(request=request, api_key="Bearer " + key)
+ # The budget must also reach the token, or the post-call increment
+ # has nothing to charge and the counter never grows.
+ assert result.user_model_max_budget == user_model_max_budget
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "over_budget,expect_refusal",
+ [(True, True), (False, False)],
+ ids=["over_budget_is_refused", "under_budget_is_served"],
+)
+async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(over_budget, expect_refusal):
+ """
+ JWT auth returns its own token instead of falling through to the
+ virtual-key budget checks, so the user's per-model budget has to be enforced
+ on that path explicitly.
+
+ The dangerous shape is not "no tracking": the post-call increment charges the
+ JWT user's counter either way, so without this check the counter grows and
+ nothing ever reads it, which looks enforced and is not.
+ """
+ from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
+ from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget
+ from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key
+ from litellm.proxy.proxy_server import model_max_budget_limiter
+
+ user_id = "jwt-user-model-budget"
+ model = "gpt-4o"
+ user_model_max_budget = {model: {"budget_limit": 1.0, "time_period": "1mo"}}
+
+ await model_max_budget_limiter.dual_cache.async_set_cache(
+ key=model_budget_spend_cache_key(
+ entity_type=Litellm_EntityType.USER,
+ entity_id=user_id,
+ budget_model=model,
+ budget_duration="1mo",
+ ),
+ value=5.0 if over_budget else 0.25,
+ ttl=600,
+ )
+
+ # The token the JWT branch builds and returns.
+ valid_token = UserAPIKeyAuth(
+ api_key=None,
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ )
+
+ if expect_refusal:
+ with pytest.raises(litellm.BudgetExceededError) as exc:
+ await _check_user_model_budget(
+ valid_token=valid_token,
+ model_max_budget_limiter=model_max_budget_limiter,
+ models=[model],
+ )
+ assert exc.value.entity_type == Litellm_EntityType.USER.value
+ else:
+ await _check_user_model_budget(
+ valid_token=valid_token,
+ model_max_budget_limiter=model_max_budget_limiter,
+ models=[model],
+ )
+
+
+def test_jwt_path_enforces_the_user_model_budget_before_returning():
+ """
+ The JWT branch returns early, so the enforcement call has to sit before that
+ return rather than in the virtual-key block. Assert on the call graph, since
+ a helper-level test passes whether or not the JWT path ever calls it.
+ """
+ import ast
+ import inspect
+ import textwrap
+
+ from litellm.proxy.auth import user_api_key_auth as auth_module
+
+ tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)))
+
+ def calls_before_each_return(node):
+ seen_check = []
+ for child in ast.walk(node):
+ if isinstance(child, ast.Call):
+ fn = child.func
+ name = getattr(fn, "id", None) or getattr(fn, "attr", None)
+ if name == "_check_user_model_budget":
+ seen_check.append(child.lineno)
+ return seen_check
+
+ check_lines = calls_before_each_return(tree)
+ assert check_lines, "_user_api_key_auth_builder never enforces the user model budget"
+
+ jwt_returns = [
+ n.lineno
+ for n in ast.walk(tree)
+ if isinstance(n, ast.Return) and isinstance(n.value, ast.Call) and getattr(n.value.func, "id", None) == "cast"
+ ]
+ assert jwt_returns, "expected the JWT branch's `return cast(UserAPIKeyAuth, valid_token)`"
+ assert any(check < jwt_return for check in check_lines for jwt_return in jwt_returns), (
+ "the user model-budget check must run before the JWT branch returns"
+ )
+
+
+def test_every_jwt_branch_carries_the_user_model_budget():
+ """
+ Each JWT branch that builds or replaces `valid_token` has to put the user's
+ model budget on it, or the enforcement call a few lines later has nothing to
+ read and silently admits the request.
+
+ The auto-register branch is the one that regressed: it REPLACES the token
+ built above it with a key-scoped one whose columns carry no user budget.
+ """
+ import ast
+ import inspect
+ import textwrap
+
+ from litellm.proxy.auth import user_api_key_auth as auth_module
+
+ tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)))
+
+ assignments = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Assign)
+ and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets)
+ ]
+ targets = {
+ t.value.id
+ for node in assignments
+ for t in node.targets
+ if isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name)
+ }
+ assert "auto_registered" in targets, (
+ f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}"
+ )
+ assert "valid_token" in targets, "the virtual-key path must carry the user's model budget"
+
+
+@pytest.mark.asyncio
+async def test_user_budget_lookup_tolerates_an_unreadable_user():
+ """
+ `get_user_object(user_id_upsert=False)` raises a bare Exception when the row
+ is simply ABSENT, which is the ordinary state for a custom-auth deployment
+ that never writes users to the proxy DB. Refusing on that exception would
+ turn "no user row" into a 4xx for every such request, and a transient DB
+ blip into a full outage.
+
+ The virtual-key path makes the same call and swallows the same exception
+ ("Unable to get user from db/cache. Setting user_obj to None"), so this is
+ the established contract, not a shortcut. There is also nothing to enforce:
+ the budget being looked up lives on the row that could not be read.
+ """
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget
+
+ prisma_client = MagicMock()
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.get_user_object",
+ new=AsyncMock(side_effect=Exception("No user table row")),
+ ):
+ budget = await _read_user_model_max_budget(
+ user_id="user-with-no-row",
+ prisma_client=prisma_client,
+ user_api_key_cache=DualCache(),
+ parent_otel_span=None,
+ proxy_logging_obj=MagicMock(),
+ )
+
+ assert budget is None
+
+
+@pytest.mark.asyncio
+async def test_user_budget_lookup_is_also_unenforced_when_the_database_is_down():
+ """
+ KNOWN LIMITATION, pinned deliberately rather than discovered later.
+
+ `get_user_object` cannot tell "row absent" from "database unreachable": the
+ absent case raises inside its own try (auth_checks.py:2177) and the handler
+ at :2213 rewrites every exception into the same
+ `ValueError("User doesn't exist in db...")`. A connection error, a query
+ timeout and a malformed row all reach us as that one type and message.
+
+ So tolerating the absent case, which the test above requires, unavoidably
+ tolerates an outage too, and a user who DOES have a per-model budget goes
+ unenforced while the DB is unreachable. This is pre-existing behaviour of
+ `get_user_object` that the virtual-key path inherits identically; it is not
+ introduced here. Distinguishing them needs a dedicated exception type for
+ the absent case and a change to both auth paths.
+ """
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget
+
+ db_down = ValueError("User doesn't exist in db. 'user_id'=u-1. Got error - Connection refused")
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.get_user_object",
+ new=AsyncMock(side_effect=db_down),
+ ):
+ budget = await _read_user_model_max_budget(
+ user_id="u-1",
+ prisma_client=MagicMock(),
+ user_api_key_cache=DualCache(),
+ parent_otel_span=None,
+ proxy_logging_obj=MagicMock(),
+ )
+
+ assert budget is None
+
+
+@pytest.mark.asyncio
+async def test_user_budget_lookup_returns_the_budget_when_the_row_reads():
+ """Positive control: the tolerance above must not be swallowing every result."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget
+
+ stored = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}
+ user_obj = MagicMock()
+ user_obj.model_max_budget = stored
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.get_user_object",
+ new=AsyncMock(return_value=user_obj),
+ ):
+ budget = await _read_user_model_max_budget(
+ user_id="user-1",
+ prisma_client=MagicMock(),
+ user_api_key_cache=DualCache(),
+ parent_otel_span=None,
+ proxy_logging_obj=MagicMock(),
+ )
+
+ assert budget == stored
+
+
+def test_zero_cost_models_skip_the_user_budget_check_on_every_path():
+ """
+ `skip_budget_checks` is computed per request for zero-cost models, and the
+ JWT branch logs "Skipping all budget checks" when it is set. Any enforcement
+ call that ignores it makes the same request behave differently depending on
+ whether the caller used a JWT or a virtual key, and makes that log a lie.
+
+ Structural rather than behavioural on purpose: the defect is a call site
+ sitting outside a guard, and driving both auth paths to a zero-cost model
+ would prove it for the two requests exercised rather than for every site.
+ """
+ import ast
+ import inspect
+ import textwrap
+
+ from litellm.proxy.auth import user_api_key_auth as auth_module
+
+ tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)))
+
+ def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool:
+ for parent in ast.walk(node):
+ if not isinstance(parent, ast.If):
+ continue
+ test = parent.test
+ is_skip_guard = (
+ isinstance(test, ast.UnaryOp)
+ and isinstance(test.op, ast.Not)
+ and isinstance(test.operand, ast.Name)
+ and test.operand.id == "skip_budget_checks"
+ )
+ if is_skip_guard and any(sub is target for sub in ast.walk(parent)):
+ return True
+ return False
+
+ calls = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget"
+ ]
+ assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}"
+
+ unguarded = [c for c in calls if not guarded_by_skip(tree, c)]
+ assert not unguarded, (
+ f"{len(unguarded)} _check_user_model_budget call(s) run even when "
+ "skip_budget_checks is set, so a zero-cost model is enforced on one auth path and not the other"
+ )
+
+
+def test_custom_auth_also_skips_budget_checks_for_zero_cost_models():
+ """
+ The custom-auth helper runs its own key, user and end-user per-model budget
+ checks. If it does not honour the zero-cost skip that the JWT and
+ virtual-key paths honour, the same free request is refused under one auth
+ method and served under the others.
+
+ Asserted structurally, on the same reasoning as the sibling test: the defect
+ is a check sitting outside a guard, and it must hold for checks added later
+ rather than only for whichever request a behavioural test happened to drive.
+ """
+ import ast
+ import inspect
+ import textwrap
+
+ from litellm.proxy.auth import user_api_key_auth as auth_module
+
+ src = textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks))
+ tree = ast.parse(src)
+
+ assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag"
+
+ budget_calls = (
+ "_check_key_model_budget_with_fallback",
+ "_check_user_model_budget",
+ "is_end_user_within_model_budget",
+ )
+
+ def guarding_ifs(target: ast.AST) -> list[ast.If]:
+ return [
+ node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is target for sub in ast.walk(node))
+ ]
+
+ def mentions_skip(node: ast.If) -> bool:
+ return any(isinstance(sub, ast.Name) and sub.id == "skip_budget_checks" for sub in ast.walk(node.test))
+
+ for call_name in budget_calls:
+ calls = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Call)
+ and (
+ (isinstance(node.func, ast.Name) and node.func.id == call_name)
+ or (isinstance(node.func, ast.Attribute) and node.func.attr == call_name)
+ )
+ ]
+ assert calls, f"{call_name} is no longer called here; update this invariant"
+ for call in calls:
+ assert any(mentions_skip(node) for node in guarding_ifs(call)), (
+ f"{call_name} runs even for a zero-cost model, so custom auth refuses "
+ "requests the JWT and virtual-key paths serve"
+ )
+
+
+def test_custom_auth_attaches_the_user_budget_even_when_it_does_not_enforce():
+ """
+ The post-call spend hook reads `user_model_max_budget` off the token, so the
+ attach has to happen whether or not THIS request was enforceable. Gating it
+ on the same condition as the check leaves the user's counter uncharged for
+ every request with no resolvable model or a zero-cost one, which is exactly
+ the untracked-spend defect this PR fixes.
+
+ Structural, because the failure is an assignment sitting inside a guard.
+ """
+ import ast
+ import inspect
+ import textwrap
+
+ from litellm.proxy.auth import user_api_key_auth as auth_module
+
+ tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks)))
+
+ attaches = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Assign)
+ and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets)
+ ]
+ assert attaches, "custom auth no longer attaches the user budget at all"
+
+ for attach in attaches:
+ enclosing_ifs = [
+ node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is attach for sub in ast.walk(node))
+ ]
+ assert not enclosing_ifs, (
+ "the user budget is attached inside a conditional, so the spend hook "
+ "cannot charge the user counter whenever that condition is false"
+ )
+
+
+def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach():
+ """
+ A JWT that maps to an existing virtual key resolves through the resolver
+ store, which builds the token from the KEY row alone and therefore carries
+ no user-level per-model budget. That branch sets `do_standard_jwt_auth =
+ False` precisely so it falls through to the shared virtual-key checks, where
+ the user row is loaded and its budget copied onto the token.
+
+ Reviewed as a bypass three times, so the two halves it depends on are pinned
+ here: the branch must not return before the shared block, and the shared
+ block must copy the user row's budget onto the token. Structural on purpose,
+ because the claim is about control flow reaching a statement, and it has to
+ hold for branches added later rather than for one mocked request.
+ """
+ import ast
+ import inspect
+ import textwrap
+
+ from litellm.proxy.auth import user_api_key_auth as auth_module
+
+ tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)))
+
+ # Half one: the shared block copies the user row's budget onto the token.
+ copies_user_row = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Assign)
+ and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets)
+ and any(isinstance(v, ast.Attribute) and v.attr == "model_max_budget" for v in ast.walk(node.value))
+ ]
+ assert copies_user_row, (
+ "nothing copies the user row's model_max_budget onto the token, so a mapped-key "
+ "JWT reaches enforcement carrying the key's columns only"
+ )
+
+ # Half two: the mapped-key branch does not return before reaching it.
+ disables_standard_auth = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Assign)
+ and any(isinstance(t, ast.Name) and t.id == "do_standard_jwt_auth" for t in node.targets)
+ and isinstance(node.value, ast.Constant)
+ and node.value.value is False
+ ]
+ assert len(disables_standard_auth) == 1, "expected exactly one mapped-key branch"
+ marker = disables_standard_auth[0]
+
+ enclosing = [
+ node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is marker for sub in node.body)
+ ]
+ assert enclosing, "could not locate the mapped-key branch body"
+
+ returns_after = [
+ node for node in ast.walk(enclosing[0]) if isinstance(node, ast.Return) and node.lineno > marker.lineno
+ ]
+ assert not returns_after, (
+ "the mapped-key branch returns before the shared virtual-key checks, so the "
+ "user's per-model budget is never attached and never enforced"
+ )
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py
index 6cc1d9e5add..28c82fdf528 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py
@@ -1202,6 +1202,8 @@ def _fake_user_api_key_auth(
model_max_budget=None,
end_user_model_max_budget=None,
end_user_id=None,
+ user_model_max_budget=None,
+ user_id=None,
token=None,
):
"""Build a minimal stand-in for ``UserAPIKeyAuth`` with just the fields
@@ -1220,6 +1222,8 @@ def _fake_user_api_key_auth(
auth.model_max_budget = model_max_budget
auth.end_user_model_max_budget = end_user_model_max_budget
auth.end_user_id = end_user_id
+ auth.user_model_max_budget = user_model_max_budget
+ auth.user_id = user_id
auth.token = token
return auth
@@ -1548,6 +1552,78 @@ async def test_summary_model_denied_when_key_over_model_budget():
assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
+async def test_summary_model_denied_when_user_over_model_budget():
+ """Internal-user per-model budget is enforced for the summary subrequest too.
+
+ This file propagates `user_api_key_user_model_max_budget` into the summary
+ subrequest's metadata, so its spend charges the user's counter. Enforcing
+ only the key and end-user scopes would let compaction increment a counter it
+ can never be refused by, which is the asymmetry this PR exists to remove.
+ """
+ import litellm
+
+ messages = _simple_messages()
+ mock_call = AsyncMock(return_value=_make_mock_response("x"))
+
+ auth = _fake_user_api_key_auth(
+ key_models=["all-proxy-models"],
+ user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}},
+ user_id="user-over-budget",
+ token="hashed-token",
+ )
+
+ limiter = MagicMock()
+ limiter.is_user_within_model_budget = AsyncMock(
+ side_effect=litellm.BudgetExceededError(
+ message="over budget", current_cost=10, max_budget=5
+ )
+ )
+
+ with (
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
+ return_value="claude-haiku-4-5",
+ ),
+ patch("litellm.token_counter", return_value=200_000),
+ patch(
+ "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
+ mock_call,
+ ),
+ patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter),
+ ):
+ result = await apply_compact_20260112(
+ model=MODEL,
+ messages=messages,
+ tools=None,
+ system=None,
+ edit_spec=_EDIT_SPEC_DEFAULT,
+ user_api_key_auth=auth,
+ )
+
+ mock_call.assert_not_awaited()
+ assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
+
+ # The limiter is a mock, so it would accept any kwargs. Pin the call shape and
+ # check it against the real method, or a rename there would keep this test
+ # green while breaking compaction in production.
+ limiter.is_user_within_model_budget.assert_awaited_once_with(
+ user_id="user-over-budget",
+ user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}},
+ model="claude-haiku-4-5",
+ )
+ import inspect
+
+ from litellm.proxy.hooks.model_max_budget_limiter import (
+ _PROXY_VirtualKeyModelMaxBudgetLimiter,
+ )
+
+ real_params = inspect.signature(
+ _PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget
+ ).parameters
+ for kwarg in ("user_id", "user_model_max_budget", "model"):
+ assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts"
+
+
async def test_summary_model_denied_when_end_user_over_model_budget():
"""End-user per-model budget is enforced for the summary subrequest too."""
import litellm
diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
index 8a036d7e62f..da51d513b39 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
@@ -2975,6 +2975,7 @@ async def test_user_info_v2_response_shape(mocker):
"updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc),
"sso_user_id": None,
"teams": ["team-a", "team-b"],
+ "model_max_budget": {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}},
}
async def mock_find_unique(*args, **kwargs):
@@ -3018,9 +3019,20 @@ async def test_user_info_v2_response_shape(mocker):
"sso_user_id",
"teams",
"object_permission",
+ "model_max_budget",
+ "model_max_budget_usage",
}
assert set(response_dict.keys()) == expected_fields
+ # The dashboard's user edit form hydrates its per-model budget rows from
+ # these two, so dropping them makes a save replace the user's budgets.
+ assert response_dict["model_max_budget"] == {
+ "gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}
+ }
+ assert response_dict["model_max_budget_usage"] == {
+ "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"}
+ }
+
# Verify teams is a list of strings (team IDs), not team objects
assert isinstance(response.teams, list)
assert all(isinstance(t, str) for t in response.teams)
@@ -4150,3 +4162,66 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker):
assert response.object_permission.mcp_tool_permissions == {
"github": ["list_issues"]
}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "model_max_budget,expected_written",
+ [
+ (
+ {"claude-opus-4-8": {"budget_limit": 200.0, "time_period": "1mo"}},
+ '{"claude-opus-4-8": {"budget_limit": 200.0, "time_period": "1mo"}}',
+ ),
+ (None, None),
+ ({}, None),
+ ],
+ ids=["supplied", "omitted", "empty"],
+)
+async def test_user_new_persists_model_max_budget(
+ monkeypatch, model_max_budget, expected_written
+):
+ """
+ /user/new used to echo model_max_budget back while writing {} to the user row,
+ so a per-model budget looked configured and was read by nothing.
+
+ The omitted/empty cases are the other half: SSO and default-key callers reach
+ generate_key_helper_fn with no budget, and writing "{}" for them would clear
+ an existing user's budgets.
+ """
+ from litellm.proxy.management_endpoints import key_management_endpoints
+
+ captured = {}
+
+ class _FakeUserRow:
+ models = []
+
+ class _FakePrisma:
+ async def insert_data(self, data, table_name):
+ if table_name == "user":
+ captured["user_data"] = dict(data)
+ return _FakeUserRow()
+ captured["key_data"] = dict(data)
+ return SimpleNamespace(
+ token=data.get("token"),
+ litellm_budget_table=None,
+ created_at=None,
+ updated_at=None,
+ )
+
+ async def get_data(self, *args, **kwargs):
+ return None
+
+ import litellm.proxy.proxy_server as proxy_server
+
+ monkeypatch.setattr(proxy_server, "prisma_client", _FakePrisma(), raising=False)
+ # model_max_budget is an enterprise feature; without this the call is rejected
+ # before it ever reaches the write this test is about.
+ monkeypatch.setattr(proxy_server, "premium_user", True, raising=False)
+
+ await key_management_endpoints.generate_key_helper_fn(
+ request_type="user",
+ user_id="u-1",
+ model_max_budget=model_max_budget,
+ )
+
+ assert captured["user_data"].get("model_max_budget") == expected_written
diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
index 069cfa01178..fff6368cfc6 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
@@ -13507,10 +13507,15 @@ async def test_info_key_fn_includes_model_max_budget_usage(monkeypatch):
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.23)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
+ # A real cache seeded at the real counter key: the spend only comes back if
+ # the endpoint computed virtual_key_spend:hashed_token_budget_test:gpt-4o:1d.
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache",
+ await _budget_cache({"virtual_key_spend:hashed_token_budget_test:gpt-4o:1d": 0.23}),
+ )
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
mock_key_info.token = test_key_token
@@ -13568,6 +13573,10 @@ async def test_info_key_fn_no_model_max_budget_skips_usage(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache",
+ mock_user_api_key_cache,
+ )
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
mock_key_info.token = test_key_token
@@ -13621,10 +13630,15 @@ async def test_info_key_fn_v2_includes_model_max_budget_usage(monkeypatch):
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.55)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
+ # A real cache seeded at the real counter key: the spend only comes back if
+ # the endpoint computed virtual_key_spend:hashed_token_v2_test:gpt-4o:7d.
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache",
+ await _budget_cache({"virtual_key_spend:hashed_token_v2_test:gpt-4o:7d": 0.55}),
+ )
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
mock_key.token = test_key_token
@@ -13680,10 +13694,15 @@ async def test_info_key_fn_budget_table_fallback(monkeypatch):
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=1.20)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
+ # A real cache seeded at the real counter key: the spend only comes back if
+ # the endpoint computed virtual_key_spend:hashed_token_budget_table_test:bedrock/anthropic.claude-opus-4:30d.
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache",
+ await _budget_cache({"virtual_key_spend:hashed_token_budget_table_test:bedrock/anthropic.claude-opus-4:30d": 1.20}),
+ )
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
mock_key_info.token = test_key_token
@@ -13748,10 +13767,15 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch):
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=2.50)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
+ # A real cache seeded at the real counter key: the spend only comes back if
+ # the endpoint computed virtual_key_spend:hashed_token_v2_bt_test:bedrock/anthropic.claude-opus-4:30d.
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache",
+ await _budget_cache({"virtual_key_spend:hashed_token_v2_bt_test:bedrock/anthropic.claude-opus-4:30d": 2.50}),
+ )
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
mock_key.token = test_key_token
@@ -13793,8 +13817,13 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch):
@pytest.mark.asyncio
-async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch):
- """Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match."""
+async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch):
+ """/key/info reads the one counter enforcement reads: the configured budget model.
+
+ It used to probe a second, provider-stripped key because the counter was
+ written under the request model instead, which is what let a key report zero
+ usage while being blocked at 429.
+ """
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import LiteLLM_VerificationToken
@@ -13808,10 +13837,15 @@ async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch):
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.75])
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
+ # A real cache seeded at the real counter key: the spend only comes back if
+ # the endpoint computed virtual_key_spend:hashed_token_prefix_test:openai/gpt-4o:7d.
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache",
+ await _budget_cache({"virtual_key_spend:hashed_token_prefix_test:openai/gpt-4o:7d": 0.75}),
+ )
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
mock_key_info.token = test_key_token
@@ -13847,7 +13881,22 @@ async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch):
assert "model_max_budget_usage" in result["info"]
usage = result["info"]["model_max_budget_usage"]
assert usage["openai/gpt-4o"]["current_spend"] == 0.75
- assert mock_user_api_key_cache.async_get_cache.await_count == 2
+
+
+async def _budget_cache(seeded):
+ """A real DualCache holding spend at the given LITERAL counter keys.
+
+ The keys are spelled out in full on purpose. Seeding via
+ model_budget_spend_cache_key would move the seed and the read together, so
+ any change to the key format would still match itself and these tests could
+ never fail, which is the exact bug they exist to catch.
+ """
+ from litellm.caching.caching import DualCache
+
+ cache = DualCache()
+ for key, spend in seeded.items():
+ await cache.async_set_cache(key, spend)
+ return cache
@pytest.mark.asyncio
@@ -13872,19 +13921,16 @@ async def test_build_model_max_budget_usage_reads_current_cache_window():
_build_model_max_budget_usage,
)
- mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.30)
+ cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-4o:30d": 0.30})
result = await _build_model_max_budget_usage(
api_key_hash="some-hash",
model_max_budget={"gpt-4o": {"budget_limit": 1.0, "time_period": "30d"}},
- user_api_key_cache=mock_user_api_key_cache,
+ user_api_key_cache=cache,
)
+ # 0.30 comes back only if the key matched virtual_key_spend:some-hash:gpt-4o:30d.
assert result["gpt-4o"]["current_spend"] == 0.30
- mock_user_api_key_cache.async_get_cache.assert_awaited_once_with(
- key="virtual_key_spend:some-hash:gpt-4o:30d"
- )
@pytest.mark.asyncio
@@ -13915,8 +13961,7 @@ async def test_build_model_max_budget_usage_skips_model_without_duration():
_build_model_max_budget_usage,
)
- mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.10)
+ cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-4o:1d": 0.10})
result = await _build_model_max_budget_usage(
api_key_hash="some-hash",
@@ -13924,11 +13969,10 @@ async def test_build_model_max_budget_usage_skips_model_without_duration():
"gpt-4o": {"budget_limit": 1.0, "time_period": "1d"},
"gpt-3.5-turbo": {"budget_limit": 0.5},
},
- user_api_key_cache=mock_user_api_key_cache,
+ user_api_key_cache=cache,
)
- assert "gpt-4o" in result
+ assert result["gpt-4o"]["current_spend"] == 0.10
assert "gpt-3.5-turbo" not in result
- assert mock_user_api_key_cache.async_get_cache.await_count == 1
@pytest.mark.asyncio
@@ -13961,8 +14005,7 @@ async def test_build_model_max_budget_usage_invalid_budget_config_skipped():
_build_model_max_budget_usage,
)
- mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.20)
+ cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-3.5-turbo:7d": 0.20})
result = await _build_model_max_budget_usage(
api_key_hash="some-hash",
@@ -13970,32 +14013,35 @@ async def test_build_model_max_budget_usage_invalid_budget_config_skipped():
"gpt-4o": {"max_budget": "not-a-number", "budget_duration": "1d"},
"gpt-3.5-turbo": {"budget_limit": 0.5, "time_period": "7d"},
},
- user_api_key_cache=mock_user_api_key_cache,
+ user_api_key_cache=cache,
)
assert "gpt-4o" not in result
- assert "gpt-3.5-turbo" in result
- assert mock_user_api_key_cache.async_get_cache.await_count == 1
+ assert result["gpt-3.5-turbo"]["current_spend"] == 0.20
@pytest.mark.asyncio
-async def test_build_model_max_budget_usage_provider_prefix_cache_fallback():
+async def test_build_model_max_budget_usage_reads_only_the_configured_model_key():
+ """One lookup, at the configured budget model.
+
+ The counter is written under the name the operator configured, so probing a
+ provider-stripped variant would read a key nothing writes.
+ """
from unittest.mock import AsyncMock
from litellm.proxy.management_endpoints.key_management_endpoints import (
_build_model_max_budget_usage,
)
- mock_user_api_key_cache = AsyncMock()
- mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.55])
+ cache = await _budget_cache({"virtual_key_spend:test-hash:openai/gpt-4o:7d": 0.55})
result = await _build_model_max_budget_usage(
api_key_hash="test-hash",
model_max_budget={"openai/gpt-4o": {"budget_limit": 2.0, "time_period": "7d"}},
- user_api_key_cache=mock_user_api_key_cache,
+ user_api_key_cache=cache,
)
+ # Seeded only under the configured name, so a provider-stripped probe reads 0.0.
assert result["openai/gpt-4o"]["current_spend"] == 0.55
- assert mock_user_api_key_cache.async_get_cache.await_count == 2
def test_list_keys_substring_matching_param_defaults_to_false():
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
index 090acf2dbb0..4c6ba23c88c 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
@@ -15,9 +15,7 @@ from fastapi import Request, UploadFile
from starlette.datastructures import FormData, Headers, QueryParams
from starlette.datastructures import UploadFile as StarletteUploadFile
-sys.path.insert(
- 0, os.path.abspath("../../..")
-) # Adds the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS,
@@ -73,9 +71,7 @@ async def test_build_request_files_from_upload_file():
upload_file = UploadFile(file=file, filename="test.txt", headers=headers)
upload_file.read = AsyncMock(return_value=file_content)
- result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(
- upload_file
- )
+ result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(upload_file)
assert result == ("test.txt", file_content, "text/plain")
# Test with Starlette UploadFile
@@ -87,9 +83,7 @@ async def test_build_request_files_from_upload_file():
)
starlette_file.read = AsyncMock(return_value=file_content)
- result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(
- starlette_file
- )
+ result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(starlette_file)
assert result == ("test2.txt", file_content, "text/plain")
@@ -275,9 +269,7 @@ async def test_non_streaming_http_request_handler_multipart_with_non_empty_parse
"""
request = MagicMock(spec=Request)
request.method = "POST"
- request.headers = Headers(
- {"content-type": "multipart/form-data; boundary=------------------------test"}
- )
+ request.headers = Headers({"content-type": "multipart/form-data; boundary=------------------------test"})
file_content = b"test file content"
file = BytesIO(file_content)
@@ -316,9 +308,7 @@ async def test_pass_through_request_failure_handler():
Critical Test: When a users pass through endpoint request fails, we must log the failure code, exception in litellm spend logs.
"""
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
- with patch(
- "litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
- ) as mock_get_client:
+ with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client:
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
) as mock_processing:
@@ -329,9 +319,7 @@ async def test_pass_through_request_failure_handler():
# Setup mock for httpx client
mock_client = MagicMock()
mock_client.client = MagicMock()
- mock_client.client.request = AsyncMock(
- side_effect=httpx.HTTPError("Request failed")
- )
+ mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed"))
mock_get_client.return_value = mock_client
# Mock headers for custom headers
@@ -364,9 +352,7 @@ async def test_pass_through_request_failure_handler():
# Verify the arguments to post_call_failure_hook
call_args = mock_proxy_logging.post_call_failure_hook.call_args[1]
assert call_args["user_api_key_dict"] == mock_user_api_key_dict
- assert isinstance(
- call_args["original_exception"], TypeError
- ) # Now expecting TypeError
+ assert isinstance(call_args["original_exception"], TypeError) # Now expecting TypeError
assert "traceback_str" in call_args
@@ -410,27 +396,14 @@ def test_is_langfuse_route():
handler = PassThroughEndpointLogging()
# Test positive cases
- assert (
- handler.is_langfuse_route("http://localhost:4000/langfuse/api/public/traces")
- is True
- )
- assert (
- handler.is_langfuse_route(
- "https://proxy.example.com/langfuse/api/public/sessions"
- )
- is True
- )
+ assert handler.is_langfuse_route("http://localhost:4000/langfuse/api/public/traces") is True
+ assert handler.is_langfuse_route("https://proxy.example.com/langfuse/api/public/sessions") is True
assert handler.is_langfuse_route("/langfuse/api/public/ingestion") is True
assert handler.is_langfuse_route("http://localhost:4000/langfuse/") is True
# Test negative cases
- assert (
- handler.is_langfuse_route("https://api.openai.com/v1/chat/completions") is False
- )
- assert (
- handler.is_langfuse_route("http://localhost:4000/anthropic/v1/messages")
- is False
- )
+ assert handler.is_langfuse_route("https://api.openai.com/v1/chat/completions") is False
+ assert handler.is_langfuse_route("http://localhost:4000/anthropic/v1/messages") is False
assert handler.is_langfuse_route("https://example.com/other") is False
assert handler.is_langfuse_route("") is False
@@ -447,17 +420,9 @@ def test_is_vertex_route_ignores_plain_predict_path_segment():
"""
handler = PassThroughEndpointLogging()
- assert (
- handler.is_vertex_route(
- "https://upstream.example.com/ml/api/v1/time-series-forecast/predict"
- )
- is False
- )
+ assert handler.is_vertex_route("https://upstream.example.com/ml/api/v1/time-series-forecast/predict") is False
assert handler.is_vertex_route("https://upstream.example.com/api/v1/search") is False
- assert (
- handler.is_vertex_route("https://upstream.example.com/predict/generateContent")
- is False
- )
+ assert handler.is_vertex_route("https://upstream.example.com/predict/generateContent") is False
assert (
handler.is_vertex_route(
@@ -483,10 +448,7 @@ def test_is_vertex_route_ignores_plain_predict_path_segment():
)
is True
)
- assert (
- handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search")
- is True
- )
+ assert handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search") is True
assert (
handler.is_vertex_route(
@@ -545,9 +507,7 @@ async def test_custom_passthrough_predict_path_logs_via_generic_handler():
mock_vertex_handler.assert_not_called()
handler._handle_logging.assert_awaited_once()
- logged_object = handler._handle_logging.call_args.kwargs[
- "standard_logging_response_object"
- ]
+ logged_object = handler._handle_logging.call_args.kwargs["standard_logging_response_object"]
assert logged_object == {"response": '{"forecast": [1, 2, 3]}'}
@@ -600,10 +560,7 @@ async def test_langfuse_passthrough_no_logging():
assert result is None
# Verify that the passthrough_logging_payload was still set (this happens before the langfuse check)
- assert (
- mock_logging_obj.model_call_details["passthrough_logging_payload"]
- == passthrough_logging_payload
- )
+ assert mock_logging_obj.model_call_details["passthrough_logging_payload"] == passthrough_logging_payload
def test_construct_target_url_with_subpath():
@@ -1051,9 +1008,7 @@ async def test_create_pass_through_route_with_cost_per_request():
# Mock the pass_through_request function to capture its call
with (
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request"
- ) as mock_pass_through,
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route"
) as mock_is_registered,
@@ -1100,10 +1055,7 @@ def test_resolve_pass_through_request_timeout_precedence():
assert resolve_pass_through_request_timeout(endpoint_timeout=800) == 800.0
with patch("litellm.proxy.proxy_server.general_settings", {}):
- assert (
- resolve_pass_through_request_timeout()
- == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS
- )
+ assert resolve_pass_through_request_timeout() == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS
def test_resolve_llm_passthrough_timeout_precedence():
@@ -1135,15 +1087,11 @@ async def test_pass_through_request_uses_resolved_timeout():
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
) as mock_get_client:
- mock_proxy_logging.pre_call_hook = AsyncMock(
- side_effect=lambda **kwargs: kwargs["data"]
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
mock_client = MagicMock()
mock_client.client = MagicMock()
- mock_client.client.request = AsyncMock(
- side_effect=httpx.HTTPError("Request failed")
- )
+ mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed"))
mock_get_client.return_value = mock_client
mock_request = MagicMock(spec=Request)
@@ -1181,9 +1129,7 @@ async def test_create_pass_through_route_forwards_timeout():
)
with (
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request"
- ) as mock_pass_through,
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route"
) as mock_is_registered,
@@ -1296,9 +1242,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body"
) as mock_get_response_body:
# Setup mock for pre_call_hook and post_call_failure_hook
- mock_proxy_logging.pre_call_hook = AsyncMock(
- return_value={"test": "data"}
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"test": "data"})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
return_value={"x-callback-test": "value"}
@@ -1308,9 +1252,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
- mock_response.aread = AsyncMock(
- return_value=b'{"success": true}'
- )
+ mock_response.aread = AsyncMock(return_value=b'{"success": true}')
mock_response.text = '{"success": true}'
mock_response.raise_for_status = MagicMock()
@@ -1330,9 +1272,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/api/endpoint"
- mock_request.body = AsyncMock(
- return_value=b'{"message": "test request"}'
- )
+ mock_request.body = AsyncMock(return_value=b'{"message": "test request"}')
mock_request.headers = Headers({})
mock_request.query_params = QueryParams({})
@@ -1411,9 +1351,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream():
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor"
) as mock_chunk_processor:
- mock_proxy_logging.pre_call_hook = AsyncMock(
- return_value={"model": "claude-3", "stream": True}
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"model": "claude-3", "stream": True})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
return_value={"x-callback-test": "value"}
@@ -1438,9 +1376,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream():
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/v1/messages"
- mock_request.body = AsyncMock(
- return_value=b'{"model": "claude-3", "stream": true}'
- )
+ mock_request.body = AsyncMock(return_value=b'{"model": "claude-3", "stream": true}')
mock_request.headers = Headers({})
mock_request.query_params = QueryParams({})
@@ -1456,9 +1392,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream():
assert async_client.send.call_args.kwargs["stream"] is True
mock_chunk_processor.assert_called_once()
- logging_obj = mock_chunk_processor.call_args.kwargs[
- "litellm_logging_obj"
- ]
+ logging_obj = mock_chunk_processor.call_args.kwargs["litellm_logging_obj"]
assert logging_obj.stream is True
assert logging_obj.model_call_details["stream"] is True
@@ -1479,9 +1413,7 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream():
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor"
) as mock_chunk_processor:
- mock_proxy_logging.pre_call_hook = AsyncMock(
- return_value={"model": "claude-3"}
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"model": "claude-3"})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
return_value={"x-callback-test": "value"}
@@ -1521,9 +1453,7 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream():
async_client.send.assert_awaited_once()
mock_chunk_processor.assert_called_once()
- logging_obj = mock_chunk_processor.call_args.kwargs[
- "litellm_logging_obj"
- ]
+ logging_obj = mock_chunk_processor.call_args.kwargs["litellm_logging_obj"]
assert logging_obj.stream is True
assert logging_obj.model_call_details["stream"] is True
@@ -1550,16 +1480,10 @@ async def test_create_pass_through_endpoint():
)
# Mock the database functions
- with patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config:
- with patch(
- "litellm.proxy.proxy_server.update_config_general_settings"
- ) as mock_update_config:
+ with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
+ with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
# Mock existing config (empty list)
- mock_get_config.return_value = ConfigFieldInfo(
- field_name="pass_through_endpoints", field_value=[]
- )
+ mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=[])
# Create test endpoint data
test_endpoint = PassThroughGenericEndpoint(
@@ -1629,12 +1553,8 @@ async def test_update_pass_through_endpoint():
)
# Mock the database functions
- with patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config:
- with patch(
- "litellm.proxy.proxy_server.update_config_general_settings"
- ) as mock_update_config:
+ with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
+ with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
# Create existing endpoint data
existing_endpoint_id = "test-endpoint-123"
existing_endpoints = [
@@ -1731,18 +1651,14 @@ async def test_create_pass_through_endpoint_auth_true_enforces_allowlist():
registry: dict = {}
with (
- patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config,
+ patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config,
patch("litellm.proxy.proxy_server.update_config_general_settings"),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
registry,
),
):
- mock_get_config.return_value = ConfigFieldInfo(
- field_name="pass_through_endpoints", field_value=[]
- )
+ mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=[])
# auth is not passed -> defaults to True on PassThroughGenericEndpoint
endpoint = PassThroughGenericEndpoint(
@@ -1757,19 +1673,12 @@ async def test_create_pass_through_endpoint_auth_true_enforces_allowlist():
)
assert any(value.get("auth") is True for value in registry.values())
- assert (
- RouteChecks.is_auth_enforced_pass_through_route(
- route="/secure-passthrough", method="POST"
- )
- is True
- )
+ assert RouteChecks.is_auth_enforced_pass_through_route(route="/secure-passthrough", method="POST") is True
post_request = MagicMock(spec=Request)
post_request.method = "POST"
- without_allowlist = UserAPIKeyAuth(
- user_id="u", allowed_routes=["llm_api_routes"]
- )
+ without_allowlist = UserAPIKeyAuth(user_id="u", allowed_routes=["llm_api_routes"])
with pytest.raises(HTTPException) as exc_info:
RouteChecks.is_virtual_key_allowed_to_call_route(
route="/secure-passthrough",
@@ -1826,9 +1735,7 @@ async def test_update_pass_through_endpoint_auth_true_enforces_allowlist():
]
with (
- patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config,
+ patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config,
patch("litellm.proxy.proxy_server.update_config_general_settings"),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
@@ -1851,19 +1758,12 @@ async def test_update_pass_through_endpoint_auth_true_enforces_allowlist():
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
)
- assert (
- RouteChecks.is_auth_enforced_pass_through_route(
- route="/edited-passthrough", method="POST"
- )
- is True
- )
+ assert RouteChecks.is_auth_enforced_pass_through_route(route="/edited-passthrough", method="POST") is True
post_request = MagicMock(spec=Request)
post_request.method = "POST"
- without_allowlist = UserAPIKeyAuth(
- user_id="u", allowed_routes=["llm_api_routes"]
- )
+ without_allowlist = UserAPIKeyAuth(user_id="u", allowed_routes=["llm_api_routes"])
with pytest.raises(HTTPException) as exc_info:
RouteChecks.is_virtual_key_allowed_to_call_route(
route="/edited-passthrough",
@@ -1905,12 +1805,8 @@ async def test_update_pass_through_endpoint_preserves_auth_false():
]
with (
- patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config,
- patch(
- "litellm.proxy.proxy_server.update_config_general_settings"
- ) as mock_update_config,
+ patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config,
+ patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
registry,
@@ -1937,12 +1833,7 @@ async def test_update_pass_through_endpoint_preserves_auth_false():
persisted = mock_update_config.call_args[1]["data"].field_value[0]
assert persisted["auth"] is False
- assert (
- RouteChecks.is_auth_enforced_pass_through_route(
- route="/public-passthrough", method="POST"
- )
- is False
- )
+ assert RouteChecks.is_auth_enforced_pass_through_route(route="/public-passthrough", method="POST") is False
@pytest.mark.asyncio
@@ -1962,9 +1853,7 @@ async def test_update_pass_through_endpoint_not_found():
)
# Mock the database functions
- with patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config:
+ with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
# Mock existing config with different endpoint
existing_endpoints = [
{
@@ -1982,9 +1871,7 @@ async def test_update_pass_through_endpoint_not_found():
)
# Create update data
- update_data = PassThroughGenericEndpoint(
- path="/test/endpoint", target="http://newapi.com/v2"
- )
+ update_data = PassThroughGenericEndpoint(path="/test/endpoint", target="http://newapi.com/v2")
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
@@ -2023,12 +1910,8 @@ async def test_delete_pass_through_endpoint():
)
# Mock the database functions
- with patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config:
- with patch(
- "litellm.proxy.proxy_server.update_config_general_settings"
- ) as mock_update_config:
+ with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
+ with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
# Create existing endpoint data
endpoint_to_delete_id = "test-endpoint-123"
other_endpoint_id = "other-endpoint-456"
@@ -2106,9 +1989,7 @@ async def test_delete_pass_through_endpoint_not_found():
)
# Mock the database functions
- with patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config:
+ with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
# Mock existing config with different endpoint
existing_endpoints = [
{
@@ -2199,14 +2080,8 @@ async def test_get_pass_through_endpoints_includes_config_and_db():
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._get_pass_through_endpoints_from_config"
) as mock_get_config:
- db_objects = [
- PassThroughGenericEndpoint(**ep, is_from_config=False)
- for ep in db_endpoints
- ]
- config_objects = [
- PassThroughGenericEndpoint(**ep, is_from_config=True)
- for ep in config_endpoints
- ]
+ db_objects = [PassThroughGenericEndpoint(**ep, is_from_config=False) for ep in db_endpoints]
+ config_objects = [PassThroughGenericEndpoint(**ep, is_from_config=True) for ep in config_endpoints]
mock_get_db.return_value = db_objects
mock_get_config.return_value = config_objects
@@ -2280,13 +2155,9 @@ async def test_delete_pass_through_endpoint_empty_list():
)
# Mock the database functions
- with patch(
- "litellm.proxy.proxy_server.get_config_general_settings"
- ) as mock_get_config:
+ with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
# Mock empty config
- mock_get_config.return_value = ConfigFieldInfo(
- field_name="pass_through_endpoints", field_value=None
- )
+ mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
@@ -2325,9 +2196,7 @@ async def test_pass_through_request_query_params_forwarding():
) as mock_get_response_body:
# Setup mock for pre_call_hook
test_body = {"name": "Azure Assistant", "model": "gpt-4o"}
- mock_proxy_logging.pre_call_hook = AsyncMock(
- return_value=test_body
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body)
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
return_value={"x-callback-test": "value"}
)
@@ -2336,9 +2205,7 @@ async def test_pass_through_request_query_params_forwarding():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
- mock_response.aread = AsyncMock(
- return_value=b'{"id": "asst_123", "object": "assistant"}'
- )
+ mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}')
mock_response.text = '{"id": "asst_123", "object": "assistant"}'
mock_response.raise_for_status = MagicMock()
@@ -2360,20 +2227,12 @@ async def test_pass_through_request_query_params_forwarding():
# Create mock request with query parameters (Azure API version)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
- mock_request.url = (
- "http://localhost:4000/azure-assistant/openai/assistants"
- )
- mock_request.body = AsyncMock(
- return_value=json.dumps(test_body).encode()
- )
- mock_request.headers = Headers(
- {"Content-Type": "application/json"}
- )
+ mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants"
+ mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode())
+ mock_request.headers = Headers({"Content-Type": "application/json"})
# Create QueryParams with api-version parameter
- mock_request.query_params = QueryParams(
- [("api-version", "2025-01-01-preview")]
- )
+ mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")])
# Create mock user API key dict
mock_user_api_key_dict = MagicMock()
@@ -2395,9 +2254,7 @@ async def test_pass_through_request_query_params_forwarding():
# The key assertion: query parameters should be preserved and passed to the HTTP handler
assert "requested_query_params" in call_kwargs
- assert call_kwargs["requested_query_params"] == {
- "api-version": "2025-01-01-preview"
- }
+ assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"}
assert call_kwargs.get("forward_multipart") is False
# Verify the target URL is correct
@@ -2447,9 +2304,7 @@ async def _run_pass_through_and_capture_wire_url(
"PassThroughEndpoint client not found in in_memory_llm_clients_cache; "
"get_async_httpx_client may not be caching this provider."
)
- cache_dict[cache_key] = SimpleNamespace(
- client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))
- )
+ cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler)))
mock_request = MagicMock(spec=Request)
mock_request.method = "GET"
@@ -2458,18 +2313,14 @@ async def _run_pass_through_and_capture_wire_url(
mock_request.body = AsyncMock(return_value=b"")
mock_proxy_logging = MagicMock()
- mock_proxy_logging.pre_call_hook = AsyncMock(
- side_effect=lambda user_api_key_dict, data, call_type: data
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
mock_proxy_logging.get_proxy_hook = MagicMock(return_value=managed_files_hook)
try:
with ExitStack() as stack:
- stack.enter_context(
- patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
- )
+ stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging))
if managed_files_hook is not None:
stack.enter_context(
patch(
@@ -2637,26 +2488,16 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter():
# Create test endpoints
endpoints = [
- PassThroughGenericEndpoint(
- id="endpoint-1", path="/api/allowed1", target="http://example.com/api1"
- ),
- PassThroughGenericEndpoint(
- id="endpoint-2", path="/api/allowed2", target="http://example.com/api2"
- ),
- PassThroughGenericEndpoint(
- id="endpoint-3", path="/api/notallowed", target="http://example.com/api3"
- ),
+ PassThroughGenericEndpoint(id="endpoint-1", path="/api/allowed1", target="http://example.com/api1"),
+ PassThroughGenericEndpoint(id="endpoint-2", path="/api/allowed2", target="http://example.com/api2"),
+ PassThroughGenericEndpoint(id="endpoint-3", path="/api/notallowed", target="http://example.com/api3"),
]
# Mock prisma client
mock_prisma_client = MagicMock()
mock_team = MagicMock()
- mock_team.metadata = {
- "allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]
- }
- mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
- return_value=mock_team
- )
+ mock_team.metadata = {"allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]}
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
@@ -2671,9 +2512,7 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter():
assert result[1].path == "/api/allowed2"
# Verify database call
- mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(
- where={"team_id": "test-team-123"}
- )
+ mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(where={"team_id": "test-team-123"})
@pytest.mark.asyncio
@@ -2691,9 +2530,7 @@ async def test_filter_endpoints_by_team_allowed_routes_team_not_found():
# Create test endpoints
endpoints = [
- PassThroughGenericEndpoint(
- id="endpoint-1", path="/api/test", target="http://example.com/api"
- ),
+ PassThroughGenericEndpoint(id="endpoint-1", path="/api/test", target="http://example.com/api"),
]
# Mock prisma client to return None (team not found)
@@ -2726,21 +2563,15 @@ async def test_filter_endpoints_by_team_allowed_routes_no_metadata():
# Create test endpoints
endpoints = [
- PassThroughGenericEndpoint(
- id="endpoint-1", path="/api/test1", target="http://example.com/api1"
- ),
- PassThroughGenericEndpoint(
- id="endpoint-2", path="/api/test2", target="http://example.com/api2"
- ),
+ PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"),
+ PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"),
]
# Mock prisma client with team that has None metadata
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = None
- mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
- return_value=mock_team
- )
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
@@ -2768,21 +2599,15 @@ async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key():
# Create test endpoints
endpoints = [
- PassThroughGenericEndpoint(
- id="endpoint-1", path="/api/test1", target="http://example.com/api1"
- ),
- PassThroughGenericEndpoint(
- id="endpoint-2", path="/api/test2", target="http://example.com/api2"
- ),
+ PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"),
+ PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"),
]
# Mock prisma client with team that has metadata but no allowed_passthrough_routes
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {"some_other_key": "some_value"}
- mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
- return_value=mock_team
- )
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
@@ -2810,21 +2635,15 @@ async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list():
# Create test endpoints
endpoints = [
- PassThroughGenericEndpoint(
- id="endpoint-1", path="/api/test1", target="http://example.com/api1"
- ),
- PassThroughGenericEndpoint(
- id="endpoint-2", path="/api/test2", target="http://example.com/api2"
- ),
+ PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"),
+ PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"),
]
# Mock prisma client with team that has empty allowed_passthrough_routes
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": []}
- mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
- return_value=mock_team
- )
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
@@ -2850,29 +2669,21 @@ async def test_filter_endpoints_by_team_allowed_routes_partial_match():
# Create test endpoints
endpoints = [
- PassThroughGenericEndpoint(
- id="endpoint-1", path="/api/openai", target="http://example.com/openai"
- ),
+ PassThroughGenericEndpoint(id="endpoint-1", path="/api/openai", target="http://example.com/openai"),
PassThroughGenericEndpoint(
id="endpoint-2",
path="/api/anthropic",
target="http://example.com/anthropic",
),
- PassThroughGenericEndpoint(
- id="endpoint-3", path="/api/azure", target="http://example.com/azure"
- ),
- PassThroughGenericEndpoint(
- id="endpoint-4", path="/api/cohere", target="http://example.com/cohere"
- ),
+ PassThroughGenericEndpoint(id="endpoint-3", path="/api/azure", target="http://example.com/azure"),
+ PassThroughGenericEndpoint(id="endpoint-4", path="/api/cohere", target="http://example.com/cohere"),
]
# Mock prisma client with team that allows only 2 routes
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]}
- mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
- return_value=mock_team
- )
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
@@ -2904,9 +2715,7 @@ async def test_bedrock_router_passthrough_metadata_initialization():
)
# Mock ProxyBaseLLMRequestProcessing to verify it's used
- with patch(
- "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing"
- ) as mock_processing_class:
+ with patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing") as mock_processing_class:
# Setup mock instance
mock_processor = MagicMock()
mock_processing_class.return_value = mock_processor
@@ -2914,12 +2723,8 @@ async def test_bedrock_router_passthrough_metadata_initialization():
# Mock successful response
mock_response = MagicMock()
mock_response.status_code = 200
- mock_response.aread = AsyncMock(
- return_value=b'{"content": [{"text": "Hello"}]}'
- )
- mock_processor.base_passthrough_process_llm_request = AsyncMock(
- return_value=mock_response
- )
+ mock_response.aread = AsyncMock(return_value=b'{"content": [{"text": "Hello"}]}')
+ mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value=mock_response)
# Create mock request with headers
mock_request = MagicMock(spec=Request)
@@ -2986,18 +2791,10 @@ async def test_bedrock_router_passthrough_metadata_initialization():
call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args[1]
# These are the critical parameters that ensure metadata is properly initialized:
- assert (
- call_kwargs["request"] == mock_request
- ), "Request must be passed for header extraction"
- assert (
- call_kwargs["user_api_key_dict"] == mock_user_api_key_dict
- ), "User API key dict needed for metadata"
- assert (
- call_kwargs["proxy_logging_obj"] == mock_proxy_logging
- ), "Logging obj needed for hooks"
- assert (
- call_kwargs["llm_router"] == mock_router
- ), "Router needed for model routing"
+ assert call_kwargs["request"] == mock_request, "Request must be passed for header extraction"
+ assert call_kwargs["user_api_key_dict"] == mock_user_api_key_dict, "User API key dict needed for metadata"
+ assert call_kwargs["proxy_logging_obj"] == mock_proxy_logging, "Logging obj needed for hooks"
+ assert call_kwargs["llm_router"] == mock_router, "Router needed for model routing"
assert call_kwargs["model"] == "my-bedrock-model", "Model name must be passed"
# Verify response was returned
@@ -3060,18 +2857,12 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata():
# Bedrock passthrough uses litellm_metadata to prevent key-level
# tags from leaking into the provider payload (GH#30629).
assert "litellm_metadata" in result, "litellm_metadata should be present in result"
- assert (
- "headers" in result["litellm_metadata"]
- ), "headers should be present in litellm_metadata"
- assert isinstance(
- result["litellm_metadata"]["headers"], dict
- ), "headers should be a dictionary"
+ assert "headers" in result["litellm_metadata"], "headers should be present in litellm_metadata"
+ assert isinstance(result["litellm_metadata"]["headers"], dict), "headers should be a dictionary"
# Verify specific headers are accessible (important for guardrails)
headers = result["litellm_metadata"]["headers"]
- assert (
- "user-agent" in headers or "User-Agent" in headers
- ), "User-Agent header should be accessible in metadata"
+ assert "user-agent" in headers or "User-Agent" in headers, "User-Agent header should be accessible in metadata"
# Also verify proxy_server_request has headers (original location)
assert "proxy_server_request" in result
@@ -3106,9 +2897,7 @@ async def test_create_pass_through_route_custom_body_url_target():
)
with (
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request"
- ) as mock_pass_through,
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route"
) as mock_is_registered,
@@ -3145,9 +2934,7 @@ async def test_create_pass_through_route_custom_body_url_target():
"retrievalQuery": {"text": "What is in the knowledge base?"},
}
- setattr(
- mock_request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, bedrock_body
- )
+ setattr(mock_request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, bedrock_body)
await endpoint_func(
request=mock_request,
@@ -3185,9 +2972,7 @@ async def test_pass_through_request_non_streaming_uses_content_for_state_raw_bod
mock_request.headers = Headers({"Content-Type": "application/json"})
mock_request.state = SimpleNamespace()
setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed)
- mock_request.body = AsyncMock(
- return_value=json.dumps(parsed_from_wire).encode("utf-8")
- )
+ mock_request.body = AsyncMock(return_value=json.dumps(parsed_from_wire).encode("utf-8"))
mock_user = MagicMock()
mock_user.api_key = "sk-test"
@@ -3260,9 +3045,7 @@ async def test_pass_through_request_streaming_uses_content_for_state_raw_body():
mock_request.headers = Headers({"Content-Type": "application/json"})
mock_request.state = SimpleNamespace()
setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed)
- mock_request.body = AsyncMock(
- return_value=json.dumps(parsed_from_wire).encode("utf-8")
- )
+ mock_request.body = AsyncMock(return_value=json.dumps(parsed_from_wire).encode("utf-8"))
mock_user = MagicMock()
mock_user.api_key = "sk-test"
@@ -3336,9 +3119,7 @@ async def test_create_pass_through_route_no_custom_body_falls_back():
)
with (
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request"
- ) as mock_pass_through,
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route"
) as mock_is_registered,
@@ -3407,32 +3188,12 @@ def test_is_registered_pass_through_route_with_custom_root():
}
with patch("litellm.proxy.utils.get_server_root_path", return_value="/proxy"):
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/proxy/api/endpoint"
- )
- is True
- )
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/api/endpoint"
- )
- is True
- )
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/proxy/api/endpoint") is True
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/api/endpoint") is True
with patch("litellm.proxy.utils.get_server_root_path", return_value="/"):
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/api/endpoint"
- )
- is True
- )
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/proxy/api/endpoint"
- )
- is False
- )
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/api/endpoint") is True
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/proxy/api/endpoint") is False
# Clean up
_registered_pass_through_routes.clear()
@@ -3464,24 +3225,18 @@ def test_get_registered_pass_through_route_with_custom_root():
with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"):
# Prefixed incoming route
- result = InitPassThroughEndpointHelpers.get_registered_pass_through_route(
- "/litellm/chat/completions"
- )
+ result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/litellm/chat/completions")
assert result is not None
assert result["target"] == "http://api.example.com/v1/chat/completions"
assert result["headers"]["Authorization"] == "Bearer token123"
# Bare incoming route (get_request_route convention)
- result = InitPassThroughEndpointHelpers.get_registered_pass_through_route(
- "/chat/completions"
- )
+ result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/chat/completions")
assert result is not None
assert result["target"] == "http://api.example.com/v1/chat/completions"
with patch("litellm.proxy.utils.get_server_root_path", return_value="/"):
- result = InitPassThroughEndpointHelpers.get_registered_pass_through_route(
- "/chat/completions"
- )
+ result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/chat/completions")
assert result is not None
assert result["target"] == "http://api.example.com/v1/chat/completions"
@@ -3535,12 +3290,7 @@ def test_db_registered_pass_through_route_bare_path_convention(
"litellm.proxy.utils.get_server_root_path",
return_value=server_root_path,
):
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- incoming_route
- )
- is should_match
- )
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route(incoming_route) is should_match
_registered_pass_through_routes.clear()
@@ -3559,25 +3309,13 @@ def test_mapped_pass_through_routes_with_server_root_path():
with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"):
# prefixed route should match mapped routes like /vertex_ai
assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/litellm/vertex_ai/v1/projects/foo"
- )
- is True
- )
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/litellm/bedrock/model/invoke"
- )
+ InitPassThroughEndpointHelpers.is_registered_pass_through_route("/litellm/vertex_ai/v1/projects/foo")
is True
)
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/litellm/bedrock/model/invoke") is True
# bare route without prefix should not match when root is set
- assert (
- InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/vertex_ai/v1/projects/foo"
- )
- is False
- )
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/vertex_ai/v1/projects/foo") is False
@pytest.mark.asyncio
@@ -3594,24 +3332,18 @@ async def test_multipart_passthrough_preserves_boundary():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = httpx.Headers({"content-type": "application/json"})
- mock_response.aread = AsyncMock(
- return_value=b'{"filename": "test.txt", "size": 17}'
- )
+ mock_response.aread = AsyncMock(return_value=b'{"filename": "test.txt", "size": 17}')
mock_response.text = '{"filename": "test.txt", "size": 17}'
async def mock_httpx_request(method, url, **kwargs):
# Verify that files parameter is passed (not json)
assert "files" in kwargs, "Files should be passed for multipart requests"
- file_parts = [
- value for name, value in kwargs["files"] if name == "file"
- ]
+ file_parts = [value for name, value in kwargs["files"] if name == "file"]
assert len(file_parts) == 1, "File field should be in files"
# Verify content-type is NOT in headers (httpx will set it with correct boundary)
headers = kwargs.get("headers", {})
- assert (
- "content-type" not in headers
- ), "content-type should be removed for multipart"
+ assert "content-type" not in headers, "content-type should be removed for multipart"
filename, content, content_type = file_parts[0]
assert filename == "test.txt"
@@ -3684,9 +3416,7 @@ def test_get_response_headers_strips_server_and_date():
"connection",
"keep-alive",
):
- assert (
- stripped not in lowered_keys
- ), f"{stripped!r} must not be forwarded by passthrough"
+ assert stripped not in lowered_keys, f"{stripped!r} must not be forwarded by passthrough"
# Application/business headers must still pass through.
lowered = {k.lower(): v for k, v in result.items()}
@@ -3724,9 +3454,7 @@ class TestStaleRouteCleanupOnReload:
)
stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True))
mock_set_env = stack.enter_context(
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.set_env_variables_in_header"
- )
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.set_env_variables_in_header")
)
mock_set_env.return_value = {}
return stack
@@ -3771,14 +3499,10 @@ class TestStaleRouteCleanupOnReload:
so the registry would hold both paths instead of only ``/b``.
"""
with self._patches():
- await initialize_pass_through_endpoints(
- [{"path": "/a", "target": "http://example.com"}]
- )
+ await initialize_pass_through_endpoints([{"path": "/a", "target": "http://example.com"}])
assert self._paths_in_registry() == ["/a"]
- await initialize_pass_through_endpoints(
- [{"path": "/b", "target": "http://example.com"}]
- )
+ await initialize_pass_through_endpoints([{"path": "/b", "target": "http://example.com"}])
assert self._paths_in_registry() == ["/b"]
@@ -3801,12 +3525,8 @@ class TestStaleRouteCleanupOnReload:
]
)
- assert InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/live-passthrough"
- )
- assert InitPassThroughEndpointHelpers.is_registered_pass_through_route(
- "/live-passthrough/some/subpath"
- )
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/live-passthrough")
+ assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/live-passthrough/some/subpath")
# Regression (LIT-3538): a pre-call guardrail block on a passthrough endpoint
@@ -3907,40 +3627,26 @@ async def _drive_pass_through_block(raised_exception):
400,
),
(
- _FastAPIHTTPException(
- status_code=400, detail={"error": "Violated moderation policy"}
- ),
+ _FastAPIHTTPException(status_code=400, detail={"error": "Violated moderation policy"}),
400,
),
],
)
-async def test_pre_call_guardrail_block_logs_warning_not_exception(
- guardrail_exception, expected_code
-):
+async def test_pre_call_guardrail_block_logs_warning_not_exception(guardrail_exception, expected_code):
status_code, logger = await _drive_pass_through_block(guardrail_exception)
assert int(status_code) == expected_code
- assert (
- logger.exception.call_count == 0
- ), "guardrail block must not be logged as an ERROR with a traceback"
- assert (
- logger.warning.call_count == 1
- ), "guardrail block must be logged once at WARNING"
+ assert logger.exception.call_count == 0, "guardrail block must not be logged as an ERROR with a traceback"
+ assert logger.warning.call_count == 1, "guardrail block must be logged once at WARNING"
@pytest.mark.asyncio
async def test_non_guardrail_exception_still_logs_with_traceback():
- status_code, logger = await _drive_pass_through_block(
- RuntimeError("upstream connection reset")
- )
+ status_code, logger = await _drive_pass_through_block(RuntimeError("upstream connection reset"))
assert int(status_code) == 500
- assert (
- logger.exception.call_count == 1
- ), "a genuine failure must still be logged via verbose_proxy_logger.exception"
- assert (
- logger.warning.call_count == 0
- ), "a genuine failure must not be downgraded to WARNING"
+ assert logger.exception.call_count == 1, "a genuine failure must still be logged via verbose_proxy_logger.exception"
+ assert logger.warning.call_count == 0, "a genuine failure must not be downgraded to WARNING"
# Regression: generic config-based passthrough (`pass_through_request`) used to
@@ -3979,9 +3685,7 @@ async def test_pass_through_request_non_streaming_upstream_error_returned_unchan
) as mock_success_handler:
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
- mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
- return_value=None
- )
+ mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_processing.get_custom_headers.return_value = {}
mock_success_handler.return_value = None
@@ -4067,9 +3771,7 @@ async def test_pass_through_request_upstream_error_failure_hook_exception_is_swa
mock_proxy_logging.post_call_failure_hook = AsyncMock(
side_effect=RuntimeError("alerting integration misconfigured")
)
- mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
- return_value=None
- )
+ mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_processing.get_custom_headers.return_value = {}
mock_success_handler.return_value = None
@@ -4119,9 +3821,7 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged(
) as mock_success_handler:
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
- mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
- return_value=None
- )
+ mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_success_handler.return_value = None
async_client = MagicMock()
@@ -4149,10 +3849,7 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged(
streamed_chunks = [chunk async for chunk in response.body_iterator]
await asyncio.sleep(0)
- streamed_bytes = b"".join(
- chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
- for chunk in streamed_chunks
- )
+ streamed_bytes = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks)
assert streamed_bytes == upstream_content
assert json.loads(streamed_bytes) == _UPSTREAM_ERROR_BODY
@@ -4198,9 +3895,7 @@ async def test_pass_through_request_non_streaming_success_unchanged():
) as mock_success_handler:
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
- mock_proxy_logging.post_call_response_headers_hook = AsyncMock(
- return_value=None
- )
+ mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_processing.get_custom_headers.return_value = {}
mock_success_handler.return_value = None
@@ -4244,9 +3939,7 @@ async def test_pass_through_request_internal_failure_still_raises_proxy_exceptio
from litellm.proxy._types import ProxyException
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
- mock_proxy_logging.pre_call_hook = AsyncMock(
- side_effect=RuntimeError("auth backend unavailable")
- )
+ mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=RuntimeError("auth backend unavailable"))
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_request = MagicMock(spec=Request)
@@ -4335,9 +4028,7 @@ def _inject_fake_passthrough_client(transport, timeout):
def _enter_relay_logging_mocks(stack, parsed_body):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
- mock_proxy_logging = stack.enter_context(
- patch("litellm.proxy.proxy_server.proxy_logging_obj")
- )
+ mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj"))
mock_proxy_logging.pre_call_hook = AsyncMock(return_value=parsed_body)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
@@ -4347,11 +4038,7 @@ def _enter_relay_logging_mocks(stack, parsed_body):
)
)
mock_success_handler.return_value = None
- stack.enter_context(
- patch.object(
- GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", new=MagicMock()
- )
- )
+ stack.enter_context(patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", new=MagicMock()))
return mock_proxy_logging, mock_success_handler
@@ -4437,10 +4124,7 @@ async def test_pass_through_request_relays_non_json_body_without_buffering():
mock_success_handler.assert_called_once()
success_kwargs = mock_success_handler.call_args.kwargs
assert success_kwargs["response_body"] is None
- assert (
- success_kwargs["url_route"]
- == "http://upstream.test/v1/messages/batches/b1/results"
- )
+ assert success_kwargs["url_route"] == "http://upstream.test/v1/messages/batches/b1/results"
finally:
cleanup()
await fake_client.aclose()
@@ -4518,9 +4202,7 @@ async def test_pass_through_request_upstream_error_body_stays_buffered():
)
try:
with ExitStack() as stack:
- mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks(
- stack, {}
- )
+ mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks(stack, {})
response = await pass_through_request(
request=_relay_client_request(),
@@ -4590,18 +4272,11 @@ async def test_pass_through_relay_client_disconnect_logs_partial_relay_warning(c
partial_relay_warnings = [
record.getMessage()
for record in caplog.records
- if record.levelno == logging.WARNING
- and _PARTIAL_RELAY_WARNING_MARKER in record.getMessage()
+ if record.levelno == logging.WARNING and _PARTIAL_RELAY_WARNING_MARKER in record.getMessage()
]
assert len(partial_relay_warnings) == 1
- assert (
- "http://upstream.test/v1/messages/batches/b1/results"
- in partial_relay_warnings[0]
- )
- assert (
- f"{len(first_chunk)} bytes were sent to the client"
- in partial_relay_warnings[0]
- )
+ assert "http://upstream.test/v1/messages/batches/b1/results" in partial_relay_warnings[0]
+ assert f"{len(first_chunk)} bytes were sent to the client" in partial_relay_warnings[0]
assert upstream_stream.closed is True
mock_success_handler.assert_called_once()
@@ -4648,10 +4323,7 @@ async def test_pass_through_relay_full_consumption_logs_no_partial_relay_warning
relayed = [chunk async for chunk in response.body_iterator]
assert b"".join(relayed) == b"".join(upstream_chunks)
- assert not any(
- _PARTIAL_RELAY_WARNING_MARKER in record.getMessage()
- for record in caplog.records
- )
+ assert not any(_PARTIAL_RELAY_WARNING_MARKER in record.getMessage() for record in caplog.records)
mock_success_handler.assert_called_once()
finally:
cleanup()
@@ -4689,9 +4361,7 @@ def _enter_upstream_usage_mocks(stack, parsed_body):
logging worker would have run so the test can await them."""
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
- mock_proxy_logging = stack.enter_context(
- patch("litellm.proxy.proxy_server.proxy_logging_obj")
- )
+ mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj"))
mock_proxy_logging.pre_call_hook = AsyncMock(return_value=parsed_body)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
@@ -4708,9 +4378,7 @@ def _enter_upstream_usage_mocks(stack, parsed_body):
return mock_proxy_logging, enqueued
-async def _run_upstream_reporting_passthrough(
- upstream_headers, status_code=200, cost_per_request=None
-):
+async def _run_upstream_reporting_passthrough(upstream_headers, status_code=200, cost_per_request=None):
"""Drive a generic pass-through against an upstream that reports its own
cost/usage. Returns (recorded standard logging payloads, proxy logging mock)."""
from litellm.proxy._types import UserAPIKeyAuth
@@ -4731,9 +4399,7 @@ async def _run_upstream_reporting_passthrough(
request=_relay_client_request(method="POST"),
target="http://internal-api.test/v1/summarize",
custom_headers={},
- user_api_key_dict=UserAPIKeyAuth(
- api_key="sk-upstream-usage", team_id="team-fil"
- ),
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"),
cost_per_request=cost_per_request,
)
for coroutine in enqueued:
@@ -4802,23 +4468,17 @@ async def test_passthrough_records_upstream_reported_cost_on_error_response():
)
mock_proxy_logging.post_call_failure_hook.assert_awaited_once()
- request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs[
- "request_data"
- ]
+ request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs["request_data"]
assert request_data["response_cost"] == 0.00021
assert request_data["combined_usage_object"] == litellm.Usage(total_tokens=930)
@pytest.mark.asyncio
async def test_passthrough_error_response_without_usage_headers_records_no_spend():
- _, mock_proxy_logging = await _run_upstream_reporting_passthrough(
- {}, status_code=500
- )
+ _, mock_proxy_logging = await _run_upstream_reporting_passthrough({}, status_code=500)
mock_proxy_logging.post_call_failure_hook.assert_awaited_once()
- request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs[
- "request_data"
- ]
+ request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs["request_data"]
assert "combined_usage_object" not in request_data
@@ -4853,14 +4513,10 @@ async def test_streaming_passthrough_records_cost_and_tokens_reported_by_upstrea
request=_relay_client_request(method="POST"),
target="http://internal-api.test/v1/summarize",
custom_headers={},
- user_api_key_dict=UserAPIKeyAuth(
- api_key="sk-upstream-usage", team_id="team-fil"
- ),
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"),
)
assert isinstance(response, StreamingResponse)
- assert [chunk async for chunk in response.body_iterator] == [
- b'data: {"delta": "hi"}\n\n'
- ]
+ assert [chunk async for chunk in response.body_iterator] == [b'data: {"delta": "hi"}\n\n']
for coroutine in enqueued:
await coroutine
finally:
@@ -4968,9 +4624,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame():
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
return_value=FakeUpstreamConnect(upstream_ws),
),
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER"
- ) as mock_worker,
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker,
):
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_success_hook = AsyncMock()
@@ -5052,9 +4706,7 @@ def _patched_websocket_passthrough_environment(upstream_ws):
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
return_value=FakeUpstreamConnect(upstream_ws),
),
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER"
- ) as mock_worker,
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker,
):
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_success_hook = AsyncMock()
@@ -5199,9 +4851,7 @@ async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rc
"abnormal": Close(1006, "connection died"),
"no_status": Close(1005, ""),
}[rcvd_close]
- upstream_ws = ClosingUpstreamWebSocket(
- ConnectionClosedError(rcvd=rcvd, sent=None, rcvd_then_sent=None)
- )
+ upstream_ws = ClosingUpstreamWebSocket(ConnectionClosedError(rcvd=rcvd, sent=None, rcvd_then_sent=None))
websocket = _client_websocket(_pending_receive)
with _patched_websocket_passthrough_environment(upstream_ws):
@@ -5276,14 +4926,15 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f
def _passthrough_kwargs_for_reservation(
- user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None
+ user_api_key_dict: UserAPIKeyAuth,
+ parsed_body: Optional[dict] = None,
+ user_defined_route: bool = False,
) -> dict:
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
- mock_request.url = (
- "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent"
- )
+ mock_request.url = "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent"
mock_request.headers = Headers({})
+ mock_request.scope = {"endpoint": _marked_pass_through_endpoint()} if user_defined_route else {}
return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=mock_request,
@@ -5352,10 +5003,7 @@ async def test_passthrough_success_reconciles_budget_reservation():
reservation = user_api_key_dict.budget_reservation
kwargs = _passthrough_kwargs_for_reservation(user_api_key_dict)
- assert (
- kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"]
- is reservation
- )
+ assert kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is reservation
increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs)
@@ -5381,9 +5029,7 @@ async def test_passthrough_body_cannot_forge_budget_reservation():
user_api_key_dict,
parsed_body={"litellm_metadata": {"user_api_key_budget_reservation": forged}},
)
- assert (
- kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None
- )
+ assert kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None
increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs)
@@ -5391,9 +5037,7 @@ async def test_passthrough_body_cannot_forge_budget_reservation():
assert increment_spend_counters.await_args.kwargs["budget_reservation"] is None
-async def _drive_streaming_pass_through(
- upstream_content_type, chunk_delay_seconds, client_asked_for_stream=True
-):
+async def _drive_streaming_pass_through(upstream_content_type, chunk_delay_seconds, client_asked_for_stream=True):
"""Drive pass_through_request against an upstream that stalls before its first byte.
``client_asked_for_stream`` picks which of pass_through_request's two streaming
@@ -5405,22 +5049,14 @@ async def _drive_streaming_pass_through(
)
with ExitStack() as stack:
- mock_proxy_logging = stack.enter_context(
- patch("litellm.proxy.proxy_server.proxy_logging_obj")
- )
+ mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj"))
mock_get_client = stack.enter_context(
- patch(
- "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
- )
- )
- mock_chunk_processor = stack.enter_context(
- patch.object(PassThroughStreamingHandler, "chunk_processor")
+ patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client")
)
+ mock_chunk_processor = stack.enter_context(patch.object(PassThroughStreamingHandler, "chunk_processor"))
mock_proxy_logging.pre_call_hook = AsyncMock(
- return_value={"model": "claude-3", "stream": True}
- if client_asked_for_stream
- else {"model": "claude-3"}
+ return_value={"model": "claude-3", "stream": True} if client_asked_for_stream else {"model": "claude-3"}
)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
@@ -5510,9 +5146,7 @@ async def test_pass_through_binary_event_stream_is_never_given_an_sse_comment():
@pytest.mark.asyncio
@pytest.mark.parametrize("configured_interval, expect_ping", [(0.05, True), (None, False)])
-async def test_pass_through_route_pings_while_the_upstream_call_is_still_running(
- configured_interval, expect_ping
-):
+async def test_pass_through_route_pings_while_the_upstream_call_is_still_running(configured_interval, expect_ping):
"""The upstream withholds its response headers until its first token, so the
whole time-to-first-token is spent inside pass_through_request with nothing on
the wire (issue #34819)."""
@@ -5542,9 +5176,7 @@ async def test_pass_through_route_pings_while_the_upstream_call_is_still_running
)
)
stack.enter_context(patch(f"{module}.pass_through_request", slow_pass_through))
- stack.enter_context(
- patch.object(litellm, "sse_keepalive_ping_interval_seconds", configured_interval)
- )
+ stack.enter_context(patch.object(litellm, "sse_keepalive_ping_interval_seconds", configured_interval))
endpoint_func = create_pass_through_route(
endpoint="/v1/messages",
@@ -5571,3 +5203,147 @@ async def test_pass_through_route_pings_while_the_upstream_call_is_still_running
assert (collected[0] == b": ping\n\n") is expect_ping
assert collected[-1] in (MESSAGE_START_SSE_FRAME, MESSAGE_START_SSE_FRAME.decode())
+
+
+def test_passthrough_carries_the_per_model_budgets():
+ """
+ Native passthrough builds its logging metadata from
+ StandardLoggingUserAPIKeyMetadata, which has no budget field, and never calls
+ add_litellm_data_to_request. Without these three keys the post-call increment
+ exits early, so a /bedrock/... request is costed but its per-model counter is
+ never written: the budget reports zero forever and enforces nothing.
+ """
+ key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}
+ user_budget = {"claude-opus-4-8": {"budget_limit": 2.0, "time_period": "1mo"}}
+ end_user_budget = {"claude-opus-4-8": {"budget_limit": 3.0, "time_period": "1d"}}
+
+ kwargs = _passthrough_kwargs_for_reservation(
+ UserAPIKeyAuth(
+ token="hash",
+ user_id="u-1",
+ model_max_budget=key_budget,
+ user_model_max_budget=user_budget,
+ end_user_model_max_budget=end_user_budget,
+ )
+ )
+
+ metadata = kwargs["litellm_params"]["metadata"]
+ assert metadata["user_api_key_model_max_budget"] == key_budget
+ assert metadata["user_api_key_user_model_max_budget"] == user_budget
+ assert metadata["user_api_key_end_user_model_max_budget"] == end_user_budget
+
+
+def test_passthrough_budget_metadata_cannot_be_forged_by_the_request_body():
+ """
+ These keys decide budget enforcement, so a caller-supplied body must not be
+ able to raise its own cap. They are set after the client metadata merge for
+ the same reason user_api_key and the parent span are.
+ """
+ key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}
+
+ kwargs = _passthrough_kwargs_for_reservation(
+ UserAPIKeyAuth(token="hash", user_id="u-1", model_max_budget=key_budget),
+ parsed_body={
+ "litellm_metadata": {
+ "user_api_key_model_max_budget": {"claude-opus-4-8": {"budget_limit": 999999.0, "time_period": "18h"}}
+ }
+ },
+ )
+
+ metadata = kwargs["litellm_params"]["metadata"]
+ assert metadata["user_api_key_model_max_budget"] == key_budget
+
+
+def _marked_pass_through_endpoint():
+ """An endpoint carrying the marker ``create_pass_through_route`` sets."""
+ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+ LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
+ )
+
+ def _endpoint(): # pragma: no cover - identity only
+ return None
+
+ setattr(_endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) # noqa: B010 # name is a module constant
+ return _endpoint
+
+
+def test_user_defined_passthrough_is_neither_tracked_nor_enforced():
+ """
+ `get_model_from_request` returns None for a user-defined pass-through on
+ purpose: the body is forwarded verbatim, so its `model` names an UPSTREAM
+ model rather than a LiteLLM-managed one, and enforcing key/team allowlists
+ against it would reject valid requests. Enforcement is therefore skipped
+ on those routes.
+
+ Attaching the budget metadata anyway would charge a counter that nothing on
+ that route can refuse, and would attribute the spend to a budget the operator
+ scoped to a LiteLLM model that merely shares the name. Tracking and
+ enforcement have to agree: both on for the built-in provider routes, both off
+ here.
+ """
+ kwargs = _passthrough_kwargs_for_reservation(
+ UserAPIKeyAuth(
+ token="hash",
+ user_id="u-1",
+ model_max_budget={"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}},
+ ),
+ user_defined_route=True,
+ )
+
+ metadata = kwargs["litellm_params"]["metadata"]
+ for field in (
+ "user_api_key_model_max_budget",
+ "user_api_key_user_model_max_budget",
+ "user_api_key_end_user_model_max_budget",
+ ):
+ assert field not in metadata, f"{field} was attached on a route that never enforces it"
+
+
+@pytest.mark.parametrize(
+ "handler_name",
+ [
+ "anthropic_proxy_route",
+ "bedrock_proxy_route",
+ "gemini_proxy_route",
+ "cohere_proxy_route",
+ "vllm_proxy_route",
+ "mistral_proxy_route",
+ ],
+)
+def test_builtin_provider_routes_do_not_carry_the_user_defined_marker(handler_name):
+ """
+ The budget metadata is attached only when the dispatched endpoint is NOT a
+ user-defined pass-through, so the built-in provider handlers must not carry
+ that marker or native provider spend would stop being tracked and enforced.
+
+ These handlers DO call `create_pass_through_route` internally, and that
+ factory sets the marker on what it returns. But the result is awaited
+ immediately rather than registered, so FastAPI puts the decorated handler in
+ `request.scope["endpoint"]`, and that is what the marker check reads. This
+ test pins the distinction between calling the factory and being dispatched as
+ its product, which is easy to misread from a grep alone.
+ """
+ from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints
+ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
+ LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
+ )
+
+ handler = getattr(llm_passthrough_endpoints, handler_name)
+ assert getattr(handler, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is False, (
+ f"{handler_name} is marked as a user-defined pass-through, so per-model budget "
+ "metadata would be skipped and native provider spend would go untracked"
+ )
+
+
+def test_the_marker_check_distinguishes_the_two_route_kinds():
+ """Positive control: the factory's product IS marked, so the check can discriminate."""
+ from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint
+ from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints
+
+ marked = MagicMock(spec=Request)
+ marked.scope = {"endpoint": _marked_pass_through_endpoint()}
+ assert request_dispatched_to_pass_through_endpoint(marked) is True
+
+ builtin = MagicMock(spec=Request)
+ builtin.scope = {"endpoint": llm_passthrough_endpoints.anthropic_proxy_route}
+ assert request_dispatched_to_pass_through_endpoint(builtin) is False
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx
index b3326df1ff8..459e3fd8c92 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx
@@ -10,6 +10,7 @@ import { Checkbox } from "@/components/ui/checkbox";
import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Separator } from "@/components/ui/separator";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
interface BulkEditUserModalProps {
open: boolean;
@@ -36,6 +37,7 @@ const BulkEditUserModal: React.FC = ({
userModels,
allowAllUsers = false,
}) => {
+ const { premiumUser } = useAuthorized();
const [loading, setLoading] = useState(false);
const [selectedTeams, setSelectedTeams] = useState([]);
const [teamBudget, setTeamBudget] = useState(null);
@@ -362,6 +364,7 @@ const BulkEditUserModal: React.FC = ({
userModels={userModels}
possibleUIRoles={possibleUIRoles}
isBulkEdit={true}
+ premiumUser={premiumUser === true}
/>
{loading && (
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx
index 0ca68cc665e..8fb94ce477e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx
@@ -1,4 +1,4 @@
-import { cleanup, screen, waitFor } from "@testing-library/react";
+import { cleanup, fireEvent, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../../../tests/test-utils";
@@ -612,6 +612,125 @@ describe("UserEditView", () => {
expect(onSubmit).not.toHaveBeenCalled();
});
+ // /user/new validates model_max_budget behind an enterprise license, so a
+ // form that re-sends what is already stored turns an unrelated edit into a
+ // 400 on a proxy without one.
+ describe("per-model budgets", () => {
+ const withStoredBudgets = {
+ ...MOCK_USER_DATA,
+ user_info: {
+ ...MOCK_USER_DATA.user_info,
+ model_max_budget: { "gpt-4": { budget_limit: 5, time_period: "30d" } },
+ },
+ };
+
+ it("should leave model_max_budget out of an edit that did not touch it", async () => {
+ const payload = await submittedPayload({ userData: withStoredBudgets, premiumUser: true });
+
+ expect(payload).not.toHaveProperty("model_max_budget");
+ });
+
+ // The proxy stores model_max_budget as a plain dict, exactly as the client
+ // sent it, and BudgetConfig documents the max_budget/budget_duration
+ // spelling. A row hydrated from the spelling the editor does not read mounts
+ // with an empty cap, and every edit re-emits ALL rows, so touching one
+ // model's budget silently deletes another's.
+ it("should keep a row stored under the BudgetConfig aliases when a sibling row is edited", async () => {
+ const onSubmit = vi.fn();
+ renderWithProviders(
+ ,
+ );
+
+ const [aliasRow, canonicalRow] = await screen.findAllByPlaceholderText("Max spend ($)");
+ expect(aliasRow).toHaveValue(5);
+
+ fireEvent.change(canonicalRow, { target: { value: "3" } });
+ await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() => {
+ expect(onSubmit).toHaveBeenCalled();
+ });
+ expect(onSubmit.mock.calls[0][0].model_max_budget).toEqual({
+ "gpt-4": { budget_limit: 5, time_period: "30d" },
+ "gpt-3.5-turbo": { budget_limit: 3, time_period: "1h" },
+ });
+ });
+
+ // The effect already re-seeds the form on a userData change, so that change
+ // does happen while this component stays mounted. The editor holds its rows
+ // in state seeded once, so without a matching re-seed the rows on screen
+ // keep describing the previously loaded user and a save overwrites theirs.
+ it("re-seeds the editor when a different user is loaded", async () => {
+ const withBudget = (limit: number, id: string) => ({
+ ...MOCK_USER_DATA,
+ user_id: id,
+ user_info: {
+ ...MOCK_USER_DATA.user_info,
+ model_max_budget: { "gpt-4": { budget_limit: limit, time_period: "1h" } },
+ },
+ });
+
+ const { rerender } = renderWithProviders(
+ ,
+ );
+ expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(5);
+
+ rerender();
+
+ expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(99);
+ });
+
+ // BulkEditUsers copies a fixed field list into its payload and never reads
+ // model_max_budget, so an editor rendered here would take input and throw
+ // it away. It also has no single stored budget to diff against, since its
+ // userData stands in for every selected user.
+ it("does not offer the editor in bulk edit, where the value would be discarded", async () => {
+ renderWithProviders(
+ ,
+ );
+
+ await screen.findByRole("button", { name: /save changes/i });
+ expect(screen.queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument();
+ });
+
+ it("should lock the editor when the proxy has no enterprise license", async () => {
+ renderWithProviders();
+
+ expect(await screen.findByPlaceholderText("Max spend ($)")).toBeDisabled();
+ });
+
+ it("should leave the editor usable when the proxy has one", async () => {
+ renderWithProviders();
+
+ expect(await screen.findByPlaceholderText("Max spend ($)")).toBeEnabled();
+ });
+ });
+
it("should send an empty-string metadata through untouched rather than as an object", async () => {
const onSubmit = vi.fn();
renderWithProviders();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx
index 5612f5cd5f6..ed0c08adf38 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx
@@ -2,6 +2,9 @@ import React, { useMemo, useState } from "react";
import { z } from "zod/v4";
import { all_admin_roles } from "@/utils/roles";
import BudgetDurationDropdown from "@/components/common_components/budget_duration_dropdown";
+import { ModelMaxBudget, ModelMaxBudgetField } from "@/components/key_team_helpers/ModelMaxBudgetEditor";
+import { modelMaxBudgetUpdate } from "@/components/key_team_helpers/modelMaxBudgetPayload";
+import { useSeededState } from "@/components/key_team_helpers/useSeededState";
import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key";
import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector";
import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions";
@@ -30,6 +33,7 @@ interface UserEditViewProps {
possibleUIRoles: Record> | null;
isBulkEdit?: boolean;
objectPermission?: ObjectPermission | null;
+ premiumUser?: boolean;
}
const MCP_SELECTION_SHAPE = z.object({
@@ -135,9 +139,14 @@ export function UserEditView({
possibleUIRoles,
isBulkEdit = false,
objectPermission,
+ premiumUser = false,
}: UserEditViewProps) {
const canEditMcpPermissions = !isBulkEdit && all_admin_roles.includes(userRole || "");
const [unlimitedBudget, setUnlimitedBudget] = useState(false);
+ const [modelMaxBudget, setModelMaxBudget] = useSeededState(
+ userData.user_id,
+ () => userData.user_info?.model_max_budget ?? {},
+ );
const schema = useMemo(() => budgetSchema(unlimitedBudget), [unlimitedBudget]);
const form = useZodForm(schema, {
defaultValues: toFormValues(userData, objectPermission, isBulkEdit, canEditMcpPermissions),
@@ -162,9 +171,11 @@ export function UserEditView({
return;
}
+ const modelBudgets = modelMaxBudgetUpdate(modelMaxBudget, userData.user_info?.model_max_budget);
onSubmit({
...values,
...("metadata" in values ? { metadata: metadata.value } : {}),
+ ...(modelBudgets !== undefined && { model_max_budget: modelBudgets }),
max_budget:
unlimitedBudget || values.max_budget === "" || values.max_budget === undefined ? null : values.max_budget,
});
@@ -282,6 +293,20 @@ export function UserEditView({
{({ id, value, onChange }) => }
+ {/* Bulk edit forwards a fixed field list and has no single stored budget to
+ diff against, so the editor would silently discard whatever was typed. */}
+ {!isBulkEdit && (
+
+ )}
+
{({ ref, value, ...control }) => (
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx
index edb0ef2802e..69254b1ffe4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx
@@ -1,4 +1,4 @@
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
import { describe, expect, it, vi, beforeEach } from "vitest";
import UserInfoView from "./user_info_view";
@@ -123,6 +123,48 @@ describe("UserInfoView add-to-team form", () => {
await user.click(await screen.findByTitle(alias));
};
+ // handleUserUpdate refreshes the local copy field by field rather than refetching,
+ // so a field it forgets reads back stale the next time the form is opened and the
+ // operator sees the save they just made apparently undone.
+ describe("per-model budgets survive a save", () => {
+ // Edit Settings lives on the details tab and is gated on write access.
+ const budgetProps = {
+ ...defaultProps,
+ userRole: "Admin",
+ initialTab: 1,
+ };
+
+ const openEditor = async (user: ReturnType) => {
+ await user.click(await screen.findByRole("button", { name: /edit settings/i }));
+ return screen.findByPlaceholderText("Max spend ($)");
+ };
+
+ beforeEach(() => {
+ mockUserGetInfoV2.mockResolvedValue({
+ ...MOCK_USER_DATA,
+ model_max_budget: { "gpt-4": { budget_limit: 5, time_period: "30d" } },
+ });
+ mockUserUpdateUserCall.mockResolvedValue({});
+ });
+
+ it("shows the saved cap, not the pre-save one, when the form is reopened", async () => {
+ const user = setup();
+ render();
+
+ fireEvent.change(await openEditor(user), { target: { value: "42" } });
+ await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() => {
+ expect(mockUserUpdateUserCall).toHaveBeenCalled();
+ });
+ expect(mockUserUpdateUserCall.mock.calls[0][1].model_max_budget).toEqual({
+ "gpt-4": { budget_limit: 42, time_period: "30d" },
+ });
+
+ expect(await openEditor(user)).toHaveValue(42);
+ });
+ });
+
it("offers only the teams the user is not already a member of", async () => {
const user = setup();
await openAddTeam(user);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx
index 32dd6f80aba..fdaa1ed3508 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx
@@ -38,6 +38,7 @@ import { ArrowLeft, CheckIcon, CopyIcon, Plus, RefreshCw, Trash2 } from "lucide-
import { toast } from "@/lib/toast";
import { getBudgetDurationLabel } from "@/components/common_components/budget_duration_dropdown";
import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import MCPServerPermissions from "@/components/permissions/MCPServerPermissions";
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
@@ -84,6 +85,7 @@ export default function UserInfoView({
initialTab = 0,
startInEditMode = false,
}: UserInfoViewProps) {
+ const { premiumUser } = useAuthorized();
const [userData, setUserData] = useState(null);
const [teamDetails, setTeamDetails] = useState([]);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
@@ -333,6 +335,7 @@ export default function UserInfoView({
max_budget: formValues.max_budget ?? userData.max_budget,
budget_duration: formValues.budget_duration ?? userData.budget_duration,
metadata: formValues.metadata ?? userData.metadata,
+ model_max_budget: formValues.model_max_budget ?? userData.model_max_budget,
object_permission: mcpEntitlement
? { ...userData.object_permission, ...mcpEntitlement }
: userData.object_permission,
@@ -391,6 +394,10 @@ export default function UserInfoView({
max_budget: userData.max_budget,
budget_duration: userData.budget_duration,
metadata: userData.metadata,
+ // Without these the per-model budget editor mounts empty and a save
+ // replaces the user's existing budgets with whatever was typed.
+ model_max_budget: userData.model_max_budget,
+ model_max_budget_usage: userData.model_max_budget_usage,
},
};
@@ -579,6 +586,7 @@ export default function UserInfoView({
userModels={userModels}
possibleUIRoles={possibleUIRoles}
objectPermission={userData.object_permission}
+ premiumUser={premiumUser === true}
/>
) : (
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx
new file mode 100644
index 00000000000..33f09092cca
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx
@@ -0,0 +1,69 @@
+import { describe, expect, it, vi } from "vitest";
+import { renderWithProviders, screen } from "../../../tests/test-utils";
+import { MODEL_MAX_BUDGET_PREMIUM_HINT, ModelMaxBudgetEditor, type ModelMaxBudget } from "./ModelMaxBudgetEditor";
+
+const STORED: ModelMaxBudget = { "gpt-4o": { budget_limit: 5, time_period: "30d" } };
+
+const renderEditor = (premiumUser: boolean, value: ModelMaxBudget = STORED) =>
+ renderWithProviders(
+ ,
+ );
+
+const addButton = () => screen.getByRole("button", { name: /Add Model Budget/i });
+
+// The proxy refuses a populated model_max_budget without an enterprise license,
+// so an editable field would only ever hand a non-premium operator a 400 after
+// they had filled the whole form in.
+describe("ModelMaxBudgetEditor without an enterprise license", () => {
+ it("locks every control on an existing row", () => {
+ renderEditor(false);
+
+ expect(screen.getByPlaceholderText("Max spend ($)")).toBeDisabled();
+ expect(addButton()).toBeDisabled();
+ });
+
+ it("still shows the budgets already stored, so they stay auditable", () => {
+ renderEditor(false);
+
+ expect(screen.getByPlaceholderText("Max spend ($)")).toHaveValue(5);
+ });
+
+ it("says why the controls are locked instead of failing silently", () => {
+ renderEditor(false);
+
+ expect(screen.getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument();
+ });
+
+ it("locks the empty state too, so no row can be started", () => {
+ renderEditor(false, {});
+
+ expect(addButton()).toBeDisabled();
+ expect(screen.getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument();
+ });
+});
+
+describe("ModelMaxBudgetEditor with an enterprise license", () => {
+ it("leaves every control usable", () => {
+ renderEditor(true);
+
+ expect(screen.getByPlaceholderText("Max spend ($)")).toBeEnabled();
+ expect(addButton()).toBeEnabled();
+ });
+
+ it("does not tell a licensed operator to upgrade", () => {
+ renderEditor(true);
+
+ expect(screen.queryByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).not.toBeInTheDocument();
+ });
+
+ it("leaves the empty state usable", () => {
+ renderEditor(true, {});
+
+ expect(addButton()).toBeEnabled();
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.test.ts b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.test.ts
new file mode 100644
index 00000000000..5bcdb29f405
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.test.ts
@@ -0,0 +1,140 @@
+import { describe, expect, it } from "vitest";
+import { entriesToModelMaxBudget, modelMaxBudgetToEntries, type ModelMaxBudget } from "./ModelMaxBudgetEditor";
+
+describe("modelMaxBudgetToEntries", () => {
+ it("hydrates an existing budget without losing its period", () => {
+ const budget: ModelMaxBudget = {
+ "claude-opus-4-8": { budget_limit: 200, time_period: "1mo" },
+ "gpt-4o": { budget_limit: 0.5, time_period: "7d" },
+ };
+ expect(modelMaxBudgetToEntries(budget)).toEqual([
+ { id: "existing-0", model: "claude-opus-4-8", budgetLimit: 200, timePeriod: "1mo", extra: {} },
+ { id: "existing-1", model: "gpt-4o", budgetLimit: 0.5, timePeriod: "7d", extra: {} },
+ ]);
+ });
+
+ it.each([
+ ["null", null],
+ ["undefined", undefined],
+ ["empty", {} as ModelMaxBudget],
+ ])("treats %s as no rows", (_label, budget) => {
+ expect(modelMaxBudgetToEntries(budget)).toEqual([]);
+ });
+});
+
+// model_max_budget is a plain dict on keys and users, stored and returned exactly
+// as the client sent it, and BudgetConfig documents the max_budget/budget_duration
+// spelling. Reading only one spelling mounts the row blank, and emitting then drops
+// it, so opening a form and saving it untouched would wipe the stored budget.
+describe("modelMaxBudgetToEntries reads either BudgetConfig spelling", () => {
+ it.each([
+ ["budget_limit/time_period", { budget_limit: 200, time_period: "1mo" }],
+ ["max_budget/budget_duration", { max_budget: 200, budget_duration: "1mo" }],
+ ])("hydrates a row stored as %s", (_label, config) => {
+ expect(modelMaxBudgetToEntries({ "claude-opus-4-8": config })).toEqual([
+ { id: "existing-0", model: "claude-opus-4-8", budgetLimit: 200, timePeriod: "1mo", extra: {} },
+ ]);
+ });
+
+ // /key/update and /budget/new both accept the limit as a string.
+ it.each([
+ ["a plain number", 0.5, 0.5],
+ ["a numeric string", "0.5", 0.5],
+ ["a trailing-zero string", "0.50", 0.5],
+ ["exponent notation", "5e-1", 0.5],
+ ["zero, which is a real cap", 0, 0],
+ ])("reads %s as the cap", (_label, stored, expected) => {
+ expect(modelMaxBudgetToEntries({ "gpt-4o": { budget_limit: stored, time_period: "1h" } })[0].budgetLimit).toBe(
+ expected,
+ );
+ });
+
+ it("leaves the cap empty rather than NaN when the stored value is not a number", () => {
+ expect(
+ modelMaxBudgetToEntries({ "gpt-4o": { budget_limit: "not a number", time_period: "1h" } })[0].budgetLimit,
+ ).toBeNull();
+ });
+
+ it("falls back to the default period rather than an empty one", () => {
+ expect(modelMaxBudgetToEntries({ "gpt-4o": { budget_limit: 1, time_period: "" } })[0].timePeriod).toBe("30d");
+ });
+});
+
+// BudgetConfig carries tpm_limit and rpm_limit too. This editor models neither, so
+// without carrying them through, editing a dollar cap silently drops a configured
+// rate limit.
+describe("fields the editor does not model", () => {
+ const WITH_RATE_LIMITS: ModelMaxBudget = {
+ "gpt-4o": { budget_limit: 5, time_period: "1h", tpm_limit: 1000, rpm_limit: 60 },
+ };
+
+ it("keeps them on the entry when hydrating", () => {
+ expect(modelMaxBudgetToEntries(WITH_RATE_LIMITS)[0].extra).toEqual({ tpm_limit: 1000, rpm_limit: 60 });
+ });
+
+ it("puts them back when the cap is edited", () => {
+ const edited = modelMaxBudgetToEntries(WITH_RATE_LIMITS).map((entry) => ({ ...entry, budgetLimit: 9 }));
+
+ expect(entriesToModelMaxBudget(edited)).toEqual({
+ "gpt-4o": { budget_limit: 9, time_period: "1h", tpm_limit: 1000, rpm_limit: 60 },
+ });
+ });
+
+ // Emitting both spellings would leave the proxy with a contradictory config.
+ it("does not re-emit the alias spelling alongside the canonical one", () => {
+ const stored: ModelMaxBudget = { "gpt-4o": { max_budget: 5, budget_duration: "1h", tpm_limit: 1000 } };
+
+ expect(entriesToModelMaxBudget(modelMaxBudgetToEntries(stored))).toEqual({
+ "gpt-4o": { budget_limit: 5, time_period: "1h", tpm_limit: 1000 },
+ });
+ });
+});
+
+describe("entriesToModelMaxBudget", () => {
+ // Two models are two independent budgets, so neither row order nor a
+ // half-filled row may change which budgets are submitted.
+ it.each([
+ ["configured first", ["gpt-4o", null] as const],
+ ["configured second", [null, "gpt-4o"] as const],
+ ])("drops a row with no model, %s", (_label, models) => {
+ const entries = models.map((model, index) => ({
+ id: String(index),
+ model,
+ budgetLimit: 1.25,
+ timePeriod: "30d",
+ extra: {},
+ }));
+ expect(entriesToModelMaxBudget(entries)).toEqual({
+ "gpt-4o": { budget_limit: 1.25, time_period: "30d" },
+ });
+ });
+
+ it("drops a row whose budget was never typed", () => {
+ expect(
+ entriesToModelMaxBudget([
+ { id: "1", model: "gpt-4o", budgetLimit: null, timePeriod: "30d", extra: {} },
+ { id: "2", model: "claude-opus-4-8", budgetLimit: 3, timePeriod: "1h", extra: {} },
+ ]),
+ ).toEqual({ "claude-opus-4-8": { budget_limit: 3, time_period: "1h" } });
+ });
+
+ it("keeps a zero budget, which is a real cap and not an empty field", () => {
+ expect(
+ entriesToModelMaxBudget([{ id: "1", model: "gpt-4o", budgetLimit: 0, timePeriod: "30d", extra: {} }]),
+ ).toEqual({
+ "gpt-4o": { budget_limit: 0, time_period: "30d" },
+ });
+ });
+
+ it("round-trips an existing budget unchanged", () => {
+ const budget: ModelMaxBudget = {
+ "claude-opus-4-8": { budget_limit: 200, time_period: "1mo" },
+ "gpt-4o": { budget_limit: 0.5, time_period: "7d" },
+ };
+ expect(entriesToModelMaxBudget(modelMaxBudgetToEntries(budget))).toEqual(budget);
+ });
+
+ it("submits nothing once the last row is removed", () => {
+ expect(entriesToModelMaxBudget([])).toEqual({});
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx
new file mode 100644
index 00000000000..c6e6a1a1495
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx
@@ -0,0 +1,233 @@
+import { SearchSelect } from "@/components/shared/SearchSelect";
+import { Field, FieldLabel } from "@/components/shared/form/field";
+import { Button } from "@/components/ui/button";
+import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Plus, X } from "lucide-react";
+import React, { useState } from "react";
+
+export interface ModelBudgetConfig {
+ budget_limit: number;
+ time_period: string;
+ /** BudgetConfig also carries tpm_limit and rpm_limit, which this editor does not model. */
+ [passthrough: string]: unknown;
+}
+
+export type ModelMaxBudget = Record;
+
+export interface ModelBudgetUsage {
+ current_spend: number;
+ budget_limit: number | null;
+ time_period: string | null;
+}
+
+interface ModelBudgetEntry {
+ id: string;
+ model: string | null;
+ budgetLimit: number | null;
+ timePeriod: string;
+ extra: Readonly>;
+}
+
+// BudgetConfig aliases budget_limit onto max_budget and time_period onto
+// budget_duration, and the proxy stores whichever spelling the client sent.
+const MODELLED_FIELDS: readonly string[] = ["budget_limit", "time_period", "max_budget", "budget_duration"];
+
+const readNumber = (raw: unknown): number | null => {
+ const value = typeof raw === "string" ? Number(raw) : raw;
+ return typeof value === "number" && Number.isFinite(value) ? value : null;
+};
+
+const readPeriod = (raw: unknown): string | null => (typeof raw === "string" && raw !== "" ? raw : null);
+
+export const MODEL_BUDGET_PERIOD_OPTIONS = [
+ { value: "1h", label: "Hourly" },
+ { value: "24h", label: "Daily" },
+ { value: "7d", label: "Weekly" },
+ { value: "30d", label: "Monthly" },
+ { value: "1mo", label: "Calendar month" },
+];
+
+const DEFAULT_PERIOD = "30d";
+
+export const entriesToModelMaxBudget = (entries: readonly ModelBudgetEntry[]): ModelMaxBudget =>
+ Object.fromEntries(
+ entries
+ .filter(
+ (entry): entry is ModelBudgetEntry & { model: string; budgetLimit: number } =>
+ entry.model !== null && entry.budgetLimit !== null,
+ )
+ .map((entry) => [
+ entry.model,
+ { ...entry.extra, budget_limit: entry.budgetLimit, time_period: entry.timePeriod },
+ ]),
+ );
+
+export const modelMaxBudgetToEntries = (budget: ModelMaxBudget | null | undefined): ModelBudgetEntry[] =>
+ Object.entries(budget ?? {}).map(([model, config], index) => ({
+ id: `existing-${index}`,
+ model,
+ budgetLimit: readNumber(config?.budget_limit) ?? readNumber(config?.max_budget),
+ timePeriod: readPeriod(config?.time_period) ?? readPeriod(config?.budget_duration) ?? DEFAULT_PERIOD,
+ extra: Object.fromEntries(Object.entries(config ?? {}).filter(([field]) => !MODELLED_FIELDS.includes(field))),
+ }));
+
+export const MODEL_MAX_BUDGET_PREMIUM_HINT = "Premium feature - Upgrade to set per-model budgets";
+
+interface ModelMaxBudgetEditorProps {
+ value: ModelMaxBudget;
+ onChange: (value: ModelMaxBudget) => void;
+ availableModels: string[];
+ /** The proxy rejects a populated model_max_budget without an enterprise license. */
+ premiumUser: boolean;
+ usage?: Record | null;
+}
+
+export function ModelMaxBudgetEditor({
+ value,
+ onChange,
+ availableModels,
+ premiumUser,
+ usage,
+}: ModelMaxBudgetEditorProps) {
+ const [entries, setEntries] = useState(() => modelMaxBudgetToEntries(value));
+
+ const emitChange = (updated: ModelBudgetEntry[]) => {
+ setEntries(updated);
+ onChange(entriesToModelMaxBudget(updated));
+ };
+
+ const addEntry = () =>
+ emitChange([
+ ...entries,
+ { id: Date.now().toString(), model: null, budgetLimit: null, timePeriod: DEFAULT_PERIOD, extra: {} },
+ ]);
+
+ const removeEntry = (id: string) => emitChange(entries.filter((entry) => entry.id !== id));
+
+ const updateEntry = (id: string, patch: Partial) =>
+ emitChange(entries.map((entry) => (entry.id === id ? { ...entry, ...patch } : entry)));
+
+ const usedModels = new Set(entries.map((entry) => entry.model).filter(Boolean));
+ const hintWhenLocked = premiumUser ? undefined : MODEL_MAX_BUDGET_PREMIUM_HINT;
+
+ const blurb = (
+
+ {premiumUser
+ ? "Cap spend per model over its own window. A budget set on the bare model name also covers the provider-prefixed spelling of that model."
+ : MODEL_MAX_BUDGET_PREMIUM_HINT}
+