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

Per-model budgets were three separate things pretending to be one. The
enforcement check, the post-call increment and the info endpoints each derived
their own cache key, so a budget could refuse traffic at 429 while /key/info
reported zero usage, and a Bedrock model id never matched a budget keyed on the
bare family name. /user/new echoed a model_max_budget back and stored an empty
dict, and nothing enforced a user-scoped per-model budget at all.

One owner now builds the counter key from the configured budget model, and
enforcement, the increment and the info endpoints all read it. Bedrock ids
resolve through the model-cost map. Auth carries the user's budget onto the
token on every branch that reaches the spend hook, including JWT and
auto-registration. Native passthrough attaches the three budget metadata keys
its StandardLoggingUserAPIKeyMetadata does not carry, so /anthropic/... and
/bedrock/... traffic is counted and capped like /v1/chat/completions.

The dashboard gains the per-model budget editor it never had, on the key create,
key edit and internal-user edit forms. It is read-only without an enterprise
license, matching the write gate the proxy already enforces, and an untouched
budget is left out of an update so an unrelated edit cannot trip that gate.

The editor hydrates from either BudgetConfig spelling, since model_max_budget is
a plain dict that the proxy stores exactly as the client sent it, and it carries
through the fields it does not model. Without both, editing one model would drop
another model row entirely and silently discard its tpm_limit and rpm_limit.

/user/info refreshes its local copy of the user field by field after a save, so
model_max_budget joins that list. Left out, a saved cap read back as the old one
when the form was reopened, and clearing the row to recover would then wipe the
value that had actually persisted.

A zero-dollar cap is the strictest limit expressible, not the absence of one,
so it is enforced rather than skipped on falsiness, spend exactly at the cap is
refused the way every sibling budget check already refuses it, and a counter
that was never written reads as zero spend rather than as unknown. The usage
endpoints read every counter in one batched lookup, so a large model_max_budget
cannot fan out into one concurrent cache call per configured model.

Every auth path honours the same zero-cost skip flag, so none of them can refuse
a free request that another serves. The custom-auth helper gains the flag it
never had, which also changes its pre-existing key and end-user checks.

The compaction summary gate checks the user scope alongside the key and end-user
ones. This file propagates all three budgets into the summary subrequest, so
enforcing only two let compaction increment a counter it could not be refused by.

Custom auth attaches the user's budget to the token unconditionally, since the
post-call spend hook reads it there: gating the attach on the same condition as
enforcement left the counter uncharged whenever the request was not itself
enforceable. An entry that will not validate is skipped rather than raised on,
so one malformed scope cannot abort every other scope's increment or turn a
config typo into a 500.

The edit forms re-seed the budget editor when a different key or user is loaded.
Its rows are seeded once and cannot re-read their own value prop, so without
this a save wrote the previously loaded record's budgets onto the current one.

Only the built-in provider pass-through routes carry the budget metadata.
get_model_from_request deliberately resolves no model for a user-defined
pass-through, since its body is forwarded verbatim and names an upstream model,
so attaching there would charge a counter nothing on that route can refuse.
This commit is contained in:
Yassin Kortam 2026-08-21 09:47:52 -07:00 committed by GitHub
parent ff02d5cfc0
commit 7da34e8aed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 3656 additions and 1041 deletions

View file

@ -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:

View file

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

View file

@ -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(

View file

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

View file

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

View file

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

View file

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

View file

@ -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,

View file

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

View file

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

View file

@ -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("<summary>x</summary>"))
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

View file

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

View file

@ -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():

View file

@ -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<BulkEditUserModalProps> = ({
userModels,
allowAllUsers = false,
}) => {
const { premiumUser } = useAuthorized();
const [loading, setLoading] = useState(false);
const [selectedTeams, setSelectedTeams] = useState<string[]>([]);
const [teamBudget, setTeamBudget] = useState<number | null>(null);
@ -362,6 +364,7 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
userModels={userModels}
possibleUIRoles={possibleUIRoles}
isBulkEdit={true}
premiumUser={premiumUser === true}
/>
{loading && (

View file

@ -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(
<UserEditView
{...defaultProps}
premiumUser={true}
onSubmit={onSubmit}
userData={{
...MOCK_USER_DATA,
user_info: {
...MOCK_USER_DATA.user_info,
model_max_budget: {
"gpt-4": { max_budget: 5, budget_duration: "30d" },
"gpt-3.5-turbo": { budget_limit: 2, time_period: "1h" },
},
},
}}
/>,
);
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(
<UserEditView {...defaultProps} premiumUser={true} userData={withBudget(5, "user-a")} />,
);
expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(5);
rerender(<UserEditView {...defaultProps} premiumUser={true} userData={withBudget(99, "user-b")} />);
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(
<UserEditView
{...defaultProps}
isBulkEdit={true}
premiumUser={true}
userData={{
...MOCK_USER_DATA,
user_info: {
...MOCK_USER_DATA.user_info,
model_max_budget: { "gpt-4": { budget_limit: 5, time_period: "1h" } },
},
}}
/>,
);
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(<UserEditView {...defaultProps} userData={withStoredBudgets} />);
expect(await screen.findByPlaceholderText("Max spend ($)")).toBeDisabled();
});
it("should leave the editor usable when the proxy has one", async () => {
renderWithProviders(<UserEditView {...defaultProps} userData={withStoredBudgets} premiumUser={true} />);
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(<UserEditView {...defaultProps} onSubmit={onSubmit} />);

View file

@ -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<string, Record<string, string>> | 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<ModelMaxBudget>(
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 }) => <BudgetDurationDropdown id={id} value={value} onChange={onChange} />}
</FormField>
{/* 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 && (
<ModelMaxBudgetField
key={userData.user_id}
premiumUser={premiumUser}
value={modelMaxBudget}
onChange={setModelMaxBudget}
availableModels={userModels}
usage={userData.user_info?.model_max_budget_usage}
hint="Cap this user's spend on individual models, each with its own reset window. Applies across every key the user holds."
/>
)}
<FormField control={form.control} name="metadata" label="Metadata">
{({ ref, value, ...control }) => (
<Textarea {...control} ref={ref} value={value ?? ""} rows={4} placeholder="Enter metadata as JSON" />

View file

@ -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<typeof userEvent.setup>) => {
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(<UserInfoView {...budgetProps} />);
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);

View file

@ -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<UserInfoV2Response | null>(null);
const [teamDetails, setTeamDetails] = useState<TeamDisplayInfo[]>([]);
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}
/>
) : (
<div className="space-y-4">

View file

@ -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(
<ModelMaxBudgetEditor
value={value}
onChange={vi.fn()}
availableModels={["gpt-4o", "claude-opus-4-8"]}
premiumUser={premiumUser}
/>,
);
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();
});
});

View file

@ -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({});
});
});

View file

@ -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<string, ModelBudgetConfig>;
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<Record<string, unknown>>;
}
// 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<string, ModelBudgetUsage> | null;
}
export function ModelMaxBudgetEditor({
value,
onChange,
availableModels,
premiumUser,
usage,
}: ModelMaxBudgetEditorProps) {
const [entries, setEntries] = useState<ModelBudgetEntry[]>(() => 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<ModelBudgetEntry>) =>
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 = (
<div className="text-xs text-muted-foreground">
{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}
</div>
);
if (entries.length === 0) {
return (
<div>
<div className="mb-2">{blurb}</div>
<Button variant="outline" size="sm" onClick={addEntry} disabled={!premiumUser} title={hintWhenLocked}>
<Plus className="w-3 h-3" />
Add Model Budget
</Button>
</div>
);
}
return (
<div className="space-y-4">
{blurb}
{entries.map((entry) => {
const modelOptions = availableModels.filter((model) => model === entry.model || !usedModels.has(model));
const spent = entry.model ? usage?.[entry.model]?.current_spend : undefined;
return (
<div key={entry.id} className="relative rounded-lg border border-border bg-muted p-4">
<button
type="button"
onClick={() => removeEntry(entry.id)}
disabled={!premiumUser}
title={hintWhenLocked}
className="absolute top-2 right-2 text-muted-foreground hover:text-destructive transition-colors p-1"
>
<X className="w-4 h-4" />
</button>
<div className="mb-3">
<label className="block text-xs font-medium text-muted-foreground mb-1">Model</label>
<SearchSelect
options={modelOptions.map((model) => ({ label: model, value: model }))}
value={entry.model ?? ""}
onValueChange={(model) => updateEntry(entry.id, { model: model === "" ? null : model })}
placeholder="Select model"
emptyText="No models found"
disabled={!premiumUser}
/>
</div>
<div className="flex gap-2 items-center">
<InputGroup className="w-40">
<InputGroupAddon>
<InputGroupText>$</InputGroupText>
</InputGroupAddon>
<InputGroupInput
type="number"
// A per-model cap is often a fraction of a cent, so a 0.01
// step would make the browser refuse the value on submit.
step="any"
min={0}
value={entry.budgetLimit ?? ""}
onChange={(event) => {
const typed = event.target.valueAsNumber;
updateEntry(entry.id, { budgetLimit: Number.isNaN(typed) ? null : typed });
}}
placeholder="Max spend ($)"
disabled={!premiumUser}
/>
</InputGroup>
<Select
items={MODEL_BUDGET_PERIOD_OPTIONS}
value={entry.timePeriod}
onValueChange={(period: string | null) => period && updateEntry(entry.id, { timePeriod: period })}
>
<SelectTrigger className="w-[150px]" disabled={!premiumUser} title={hintWhenLocked}>
<SelectValue />
</SelectTrigger>
<SelectContent>
{MODEL_BUDGET_PERIOD_OPTIONS.map((option) => (
<SelectItem key={option.value} value={option.value}>
{option.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
{spent !== undefined && (
<div className="text-[11px] text-muted-foreground mt-2 ml-1">
Current window spend: ${spent}
{entry.budgetLimit !== null && ` of $${entry.budgetLimit}`}
</div>
)}
</div>
);
})}
<Button variant="outline" size="sm" onClick={addEntry} disabled={!premiumUser} title={hintWhenLocked}>
<Plus className="w-3 h-3" />
Add Model Budget
</Button>
</div>
);
}
interface ModelMaxBudgetFieldProps extends ModelMaxBudgetEditorProps {
hint: string;
}
/** The editor with its label, so every form that offers it presents it the same way. */
export function ModelMaxBudgetField({ hint, ...editorProps }: ModelMaxBudgetFieldProps) {
return (
<Field>
<FieldLabel>
<span title={hint}>Per-Model Budgets</span>
</FieldLabel>
<ModelMaxBudgetEditor {...editorProps} />
</Field>
);
}

View file

@ -2,6 +2,7 @@ import { Setter } from "@/types";
import { useEffect, useState } from "react";
import { keyListCall, Member, Organization } from "../networking";
import type { ObjectPermission } from "../object_permission_types";
import type { ModelBudgetUsage, ModelMaxBudget } from "./ModelMaxBudgetEditor";
export interface Team {
team_id: string;
@ -51,7 +52,8 @@ export interface KeyResponse {
key_type: string | null;
permissions: Record<string, unknown>;
model_spend: Record<string, number>;
model_max_budget: Record<string, number>;
model_max_budget: ModelMaxBudget;
model_max_budget_usage?: Record<string, ModelBudgetUsage> | null;
soft_budget_cooldown: boolean;
blocked: boolean;
litellm_budget_table: Record<string, unknown>;

View file

@ -0,0 +1,71 @@
import { describe, expect, it } from "vitest";
import { modelMaxBudgetUpdate } from "./modelMaxBudgetPayload";
const GPT_4O = { "gpt-4o": { budget_limit: 5, time_period: "30d" } };
describe("modelMaxBudgetUpdate", () => {
it("sends the edited budgets when a row is added to a key that had none", () => {
expect(modelMaxBudgetUpdate(GPT_4O, {})).toEqual(GPT_4O);
});
// Omitting the key would leave the deleted row enforcing, so a cleared editor
// has to send an explicit empty map.
it("sends {} when the last row is removed from a budget that was stored", () => {
expect(modelMaxBudgetUpdate({}, GPT_4O)).toEqual({});
});
it.each([
["null", null],
["undefined", undefined],
["empty", {}],
])("omits the key entirely when nothing was stored (%s) and nothing was entered", (_label, stored) => {
expect(modelMaxBudgetUpdate({}, stored)).toBeUndefined();
});
// Re-sending an unchanged budget is not merely wasteful: /key/update validates
// the field whenever it is present and rejects it without an enterprise
// license, so editing an unrelated field would start failing with a 400.
describe("omits an unchanged budget so an unrelated edit does not trip the license check", () => {
it("when the stored value is byte-identical", () => {
expect(modelMaxBudgetUpdate(GPT_4O, { "gpt-4o": { budget_limit: 5, time_period: "30d" } })).toBeUndefined();
});
it("when the proxy stored it under the BudgetConfig aliases", () => {
expect(modelMaxBudgetUpdate(GPT_4O, { "gpt-4o": { max_budget: 5, budget_duration: "30d" } })).toBeUndefined();
});
it("when a CRUD endpoint stored the limit as a string", () => {
expect(modelMaxBudgetUpdate(GPT_4O, { "gpt-4o": { budget_limit: "5", time_period: "30d" } })).toBeUndefined();
});
it("when only the key order differs", () => {
const edited = {
"gpt-4o": { budget_limit: 5, time_period: "30d" },
"claude-opus-4-8": { budget_limit: 1, time_period: "1h" },
};
expect(
modelMaxBudgetUpdate(edited, {
"claude-opus-4-8": { budget_limit: 1, time_period: "1h" },
"gpt-4o": { budget_limit: 5, time_period: "30d" },
}),
).toBeUndefined();
});
});
describe("sends the edited budgets whenever anything actually changed", () => {
it.each([
["the limit", { "gpt-4o": { budget_limit: 6, time_period: "30d" } }],
["the period", { "gpt-4o": { budget_limit: 5, time_period: "7d" } }],
["the model", { "gpt-4o-mini": { budget_limit: 5, time_period: "30d" } }],
["an added model", { ...GPT_4O, "claude-opus-4-8": { budget_limit: 1, time_period: "1h" } }],
])("%s", (_label, stored) => {
expect(modelMaxBudgetUpdate(GPT_4O, stored)).toEqual(GPT_4O);
});
// 0 is a real cap, so it must not compare equal to "no limit stored".
it("a zero cap replacing a stored budget with no limit at all", () => {
const zeroed = { "gpt-4o": { budget_limit: 0, time_period: "30d" } };
expect(modelMaxBudgetUpdate(zeroed, { "gpt-4o": { time_period: "30d" } })).toEqual(zeroed);
});
});
});

View file

@ -0,0 +1,44 @@
import type { ModelMaxBudget } from "./ModelMaxBudgetEditor";
/**
* A stored budget can carry either spelling: the proxy's BudgetConfig aliases
* `budget_limit`/`time_period` onto `max_budget`/`budget_duration`, and its CRUD
* endpoints accept the limit as a string.
*/
export type StoredModelMaxBudget = Record<
string,
| {
budget_limit?: number | string | null;
time_period?: string | null;
max_budget?: number | string | null;
budget_duration?: string | null;
}
| null
| undefined
>;
const canonical = (budget: StoredModelMaxBudget | null | undefined): string =>
JSON.stringify(
Object.entries(budget ?? {})
.map(([model, config]) => [
model,
Number(config?.budget_limit ?? config?.max_budget ?? NaN),
config?.time_period ?? config?.budget_duration ?? null,
])
.sort((left, right) => String(left[0]).localeCompare(String(right[0]))),
);
/**
* What to send for `model_max_budget` on an update, or undefined to omit the key.
*
* Omitting an unchanged budget matters beyond saving bytes: the write is gated on
* an enterprise license, so re-sending what is already stored makes an unrelated
* edit fail with a 400 on a proxy without one.
*
* Clearing the last row still has to send `{}`: omitting the key leaves the stored
* budgets in place, so the row the operator just deleted would keep enforcing.
*/
export const modelMaxBudgetUpdate = (
edited: ModelMaxBudget,
stored: StoredModelMaxBudget | null | undefined,
): ModelMaxBudget | undefined => (canonical(edited) === canonical(stored) ? undefined : edited);

View file

@ -0,0 +1,34 @@
import type { ModelMaxBudget } from "./ModelMaxBudgetEditor";
import { modelMaxBudgetUpdate, type StoredModelMaxBudget } from "./modelMaxBudgetPayload";
import { useSeededState } from "./useSeededState";
interface ModelMaxBudgetField<TValues> {
readonly value: ModelMaxBudget;
readonly setValue: (next: ModelMaxBudget) => void;
/** Adds `model_max_budget` to the payload only when it actually changed. */
readonly applyTo: (values: TValues) => void;
}
/**
* The seeding and submit halves of the per-model budget field, kept together
* because they share one invariant: both compare against the SAME stored value.
* Seeding from one record while diffing against another is how an edit form
* ends up writing the previously loaded key's budgets onto the current one.
*/
export function useModelMaxBudgetField<TValues extends { model_max_budget?: ModelMaxBudget }>(
identity: unknown,
stored: StoredModelMaxBudget | null | undefined,
): ModelMaxBudgetField<TValues> {
const [value, setValue] = useSeededState<ModelMaxBudget>(identity, () => (stored ?? {}) as ModelMaxBudget);
return {
value,
setValue,
applyTo: (values: TValues) => {
const update = modelMaxBudgetUpdate(value, stored);
if (update !== undefined) {
values.model_max_budget = update;
}
},
};
}

View file

@ -0,0 +1,26 @@
import { useState } from "react";
/**
* State seeded from `seed`, re-seeded whenever `identity` changes.
*
* The per-model budget editor holds its rows in state seeded once on mount, so
* it cannot re-read its `value` prop: that prop is the state it feeds, and
* re-hydrating on every change would wipe a half-typed row. Loading a different
* key or user therefore has to re-seed here, or the rows on screen keep
* describing the previously loaded one and a save overwrites their budgets.
*
* Adjusting state during render is React's documented way to reset state on a
* prop change. An effect would paint one frame with the stale value first, and
* the dashboard's lint rules reject synchronous setState inside an effect.
*/
export function useSeededState<T>(identity: unknown, seed: () => T): [T, (next: T) => void] {
const [value, setValue] = useState<T>(seed);
const [seededFrom, setSeededFrom] = useState(identity);
if (seededFrom !== identity) {
setSeededFrom(identity);
setValue(seed());
}
return [value, setValue];
}

View file

@ -58,6 +58,7 @@ import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } fro
import { Team } from "./key_team_helpers/key_list";
import { EmailEventSettingsResponse, EmailEventSettingsUpdateRequest } from "./email_events/types";
import type { SkillRegisterRequest } from "./claude_code_plugins/types";
import type { ModelBudgetUsage, ModelMaxBudget } from "./key_team_helpers/ModelMaxBudgetEditor";
import type { ObjectPermission } from "./object_permission_types";
import { jsonFields } from "./common_components/check_openapi_schema";
import type { MCPUserEnvVarsStatus } from "./mcp_tools/types";
@ -1045,6 +1046,8 @@ export interface UserInfoV2Response {
sso_user_id: string | null;
teams: string[];
object_permission?: ObjectPermission | null;
model_max_budget?: ModelMaxBudget | null;
model_max_budget_usage?: Record<string, ModelBudgetUsage> | null;
}
/**

View file

@ -16,6 +16,7 @@ const baseInput: KeyCreateInput = {
budgetLimits: [],
tagRateLimits: [],
budgetFallbacks: {},
modelMaxBudget: {},
};
const build = (formValues: Record<string, unknown>, overrides: Partial<KeyCreateInput> = {}): KeyPayloadResult =>
@ -554,3 +555,27 @@ describe("endpoint", () => {
expect(result.kind === "ok" && result.endpoint).toBe(endpoint);
});
});
describe("model_max_budget", () => {
it("sends the per-model budgets in the shape the API stores", () => {
const payload = payloadOf(
build(
{},
{
modelMaxBudget: {
"claude-opus-4-8": { budget_limit: 200, time_period: "1mo" },
"gpt-4o": { budget_limit: 0.5, time_period: "30d" },
},
},
),
);
expect(payload.model_max_budget).toEqual({
"claude-opus-4-8": { budget_limit: 200, time_period: "1mo" },
"gpt-4o": { budget_limit: 0.5, time_period: "30d" },
});
});
it("omits model_max_budget entirely when no model budget is set", () => {
expect(payloadOf(build({}))).not.toHaveProperty("model_max_budget");
});
});

View file

@ -2,6 +2,7 @@ import { mapDisplayToInternalNames } from "../callback_info_helpers";
import { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown";
import type { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion";
import type { BudgetWindowEntry } from "../key_team_helpers/BudgetWindowsEditor";
import type { ModelMaxBudget } from "../key_team_helpers/ModelMaxBudgetEditor";
import { tagRowsToLimits, type TagRateLimitEntry } from "../key_team_helpers/TagRateLimitEditor";
export interface KeyLoggingSetting {
@ -28,6 +29,7 @@ export interface KeyCreateInput {
readonly budgetLimits: BudgetWindowEntry[];
readonly tagRateLimits: TagRateLimitEntry[];
readonly budgetFallbacks: Record<string, string[]>;
readonly modelMaxBudget: ModelMaxBudget;
}
export type KeyPayloadResult =
@ -206,6 +208,7 @@ export const buildKeyCreatePayload = (input: KeyCreateInput): KeyPayloadResult =
...(validWindows.length > 0 && { budget_limits: validWindows }),
...(Object.keys(tag_rpm_limit).length > 0 && { tag_rpm_limit }),
...(Object.keys(input.budgetFallbacks).length > 0 && { budget_fallbacks: input.budgetFallbacks }),
...(Object.keys(input.modelMaxBudget).length > 0 && { model_max_budget: input.modelMaxBudget }),
...(values.budget_duration === NEVER_RESETS_BUDGET_DURATION && { budget_duration: null }),
},
};

View file

@ -61,6 +61,7 @@ import ProjectDropdown from "../common_components/ProjectDropdown";
import { CreateUserButton } from "../CreateUserButton";
import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor";
import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
import { ModelMaxBudget, ModelMaxBudgetEditor } from "../key_team_helpers/ModelMaxBudgetEditor";
import { TagRateLimitEditor, TagRateLimitEntry } from "../key_team_helpers/TagRateLimitEditor";
import {
excludeProxyWideSentinel,
@ -280,6 +281,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
const routerSettingsRef = useRef<RouterSettingsAccordionRef>(null);
const [budgetLimits, setBudgetLimits] = useState<BudgetWindowEntry[]>([]);
const [modelMaxBudget, setModelMaxBudget] = useState<ModelMaxBudget>({});
const [tagRateLimits, setTagRateLimits] = useState<TagRateLimitEntry[]>([]);
const [budgetFallbacks, setBudgetFallbacks] = useState<Record<string, string[]>>({});
const [budgetFallbacksKey, setBudgetFallbacksKey] = useState<number>(0);
@ -449,6 +451,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
modelAliases,
routerSettings: routerSettingsRef.current?.getValue() ?? routerSettings,
budgetLimits,
modelMaxBudget,
tagRateLimits,
budgetFallbacks,
};
@ -1063,6 +1066,22 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
</FieldLabel>
<BudgetWindowsEditor value={budgetLimits} onChange={setBudgetLimits} />
</Field>
<Field className="mt-4">
<FieldLabel>
<span>
Per-Model Budgets{" "}
<SimpleTooltip content="Cap spend on individual models, each with its own reset window. Enforced across every request this key makes; usage is reported on the key's info page.">
<Info className="ml-1 inline size-3.5 align-text-bottom" />
</SimpleTooltip>
</span>
</FieldLabel>
<ModelMaxBudgetEditor
value={modelMaxBudget}
onChange={setModelMaxBudget}
availableModels={modelsToPick}
premiumUser={premiumUser === true}
/>
</Field>
<Field className="mt-4">
<FieldLabel>
<span>

View file

@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../tests/test-utils";
import { KeyResponse } from "../key_team_helpers/key_list";
import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "../key_team_helpers/ModelMaxBudgetEditor";
import {
getPassThroughEndpointsCall,
getPoliciesList,
@ -1205,6 +1206,94 @@ describe("KeyEditView", () => {
});
});
describe("per-model budgets", () => {
const keyDataWithBudgets = {
...MOCK_KEY_DATA,
model_max_budget: { "gpt-4": { budget_limit: 5, time_period: "30d" } },
};
const renderWith = (premiumUser: boolean) => {
const onSubmit = vi.fn().mockResolvedValue(undefined);
renderWithProviders(
<KeyEditView
keyData={keyDataWithBudgets}
onCancel={() => {}}
onSubmit={onSubmit}
accessToken={"test-token"}
userID={"test-user"}
userRole={"admin"}
premiumUser={premiumUser}
/>,
);
return onSubmit;
};
// Same hazard as the user edit form: keyData changes while this component
// stays mounted (the effect re-seeds the form for exactly that reason), and
// the editor's rows live in state seeded once.
it("re-seeds the editor when a different key is loaded", async () => {
const withBudget = (limit: number, token: string) => ({
...MOCK_KEY_DATA,
token,
model_max_budget: { "gpt-4": { budget_limit: limit, time_period: "1h" } },
});
const { rerender } = renderWithProviders(
<KeyEditView
keyData={withBudget(5, "tok-a")}
onCancel={() => {}}
onSubmit={vi.fn().mockResolvedValue(undefined)}
accessToken={"test-token"}
userID={"test-user"}
userRole={"admin"}
premiumUser={true}
/>,
);
expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(5);
rerender(
<KeyEditView
keyData={withBudget(99, "tok-b")}
onCancel={() => {}}
onSubmit={vi.fn().mockResolvedValue(undefined)}
accessToken={"test-token"}
userID={"test-user"}
userRole={"admin"}
premiumUser={true}
/>,
);
expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(99);
});
it("should say why the editor is locked when the proxy has no enterprise license", async () => {
renderWith(false);
expect(await screen.findByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument();
});
it("should leave the editor usable when the proxy has one", async () => {
renderWith(true);
expect(await screen.findByText(/Cap spend per model over its own window/)).toBeInTheDocument();
expect(screen.queryByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).not.toBeInTheDocument();
});
// /key/update validates model_max_budget whenever the field is present and
// rejects it without a license, so re-sending an untouched budget would turn
// every unrelated edit into a 400.
it("should leave model_max_budget out of an edit that did not touch it", async () => {
const onSubmit = renderWith(true);
await userEvent.click(await screen.findByRole("button", { name: /save changes/i }));
await waitFor(() => {
expect(onSubmit).toHaveBeenCalled();
});
expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("model_max_budget");
});
});
it("should display 'AI APIs' label for the llm_api key type option", async () => {
const keyDataWithLlmApiRoutes = {
...MOCK_KEY_DATA,

View file

@ -42,6 +42,8 @@ import {
toSubmittedValues,
} from "./keyEditFormValues";
import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor";
import { ModelMaxBudgetField } from "../key_team_helpers/ModelMaxBudgetEditor";
import { useModelMaxBudgetField } from "../key_team_helpers/useModelMaxBudgetField";
import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
import {
TagRateLimitEditor,
@ -117,6 +119,7 @@ export function KeyEditView({
const [budgetFallbacks, setBudgetFallbacks] = useState<Record<string, string[]>>(
keyData.budget_fallbacks && typeof keyData.budget_fallbacks === "object" ? keyData.budget_fallbacks : {},
);
const modelBudget = useModelMaxBudgetField(keyData.token, keyData.model_max_budget);
const routerSettingsRef = useRef<RouterSettingsAccordionRef>(null);
const keyTypeFieldId = React.useId();
const projectFieldId = React.useId();
@ -284,6 +287,8 @@ export function KeyEditView({
values.budget_fallbacks = {};
}
modelBudget.applyTo(values);
const routerSettings = routerSettingsUpdate(
routerSettingsRef.current?.getValue()?.router_settings,
keyData.router_settings,
@ -444,6 +449,16 @@ export function KeyEditView({
<BudgetWindowsEditor value={budgetLimits} onChange={setBudgetLimits} />
</Field>
<ModelMaxBudgetField
key={keyData.token}
premiumUser={premiumUser}
value={modelBudget.value}
onChange={modelBudget.setValue}
availableModels={availableModels}
usage={keyData.model_max_budget_usage}
hint="Cap spend on individual models, each with its own reset window. Enforced across every request this key makes."
/>
<Field>
<FieldLabel>
{labelWithHint(

View file

@ -36086,6 +36086,10 @@ export interface components {
user_id?: string | null;
/** User Max Budget */
user_max_budget?: number | null;
/** User Model Max Budget */
user_model_max_budget?: {
[key: string]: unknown;
} | null;
user_role?: components["schemas"]["LitellmUserRoles"] | null;
/** User Rpm Limit */
user_rpm_limit?: number | null;
@ -36190,6 +36194,14 @@ export interface components {
metadata?: {
[key: string]: unknown;
} | null;
/** Model Max Budget */
model_max_budget?: {
[key: string]: unknown;
} | null;
/** Model Max Budget Usage */
model_max_budget_usage?: {
[key: string]: unknown;
} | null;
/**
* Models
* @default []