Merge pull request #41636 from BerriAI/litellm_per_key_end_user_default_budget

feat(proxy): per-key default budget for dynamically created customers
This commit is contained in:
Yassin Kortam 2026-09-18 11:36:59 -07:00 committed by GitHub
commit ca79c393d5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 1354 additions and 60 deletions

View file

@ -1219,6 +1219,7 @@ class KeyRequestBase(GenerateRequestBase):
default_estimated_output_tokens: PositiveInt | None = None
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None
budget_id: str | None = None
end_user_budget_id: str | None = None
tags: list[str] | None = None
disable_global_guardrails: bool | None = None
enable_prompt_caching: bool | None = None
@ -4734,6 +4735,7 @@ LiteLLM_ManagementEndpoint_MetadataFields: Final = [
"enforced_file_expires_after",
"throttle_on_budget_exceeded",
"enable_prompt_caching",
"end_user_budget_id",
]
LiteLLM_ManagementEndpoint_MetadataFields_Premium: Final = [

View file

@ -1353,29 +1353,44 @@ def get_actual_routes(allowed_routes: list) -> list:
return actual_routes
KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id"
def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None:
"""The default budget a key assigns to end users that carry no budget of their own."""
if key_metadata is None:
return None
budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD)
return budget_id if isinstance(budget_id, str) and budget_id != "" else None
async def get_default_end_user_budget(
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None = None,
budget_id: str | None = None,
) -> LiteLLM_BudgetTable | None:
"""
Fetches the default end user budget from the database if litellm.max_end_user_budget_id is configured.
Fetches the default end user budget from the database.
This budget is applied to end users who don't have an explicit budget_id set.
Results are cached for performance.
``budget_id`` selects the budget row; when omitted the proxy-wide
``litellm.max_end_user_budget_id`` is used. This budget is applied to end
users who don't have an explicit budget_id set. Results are cached for performance.
Args:
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving budget data
parent_otel_span: Optional OpenTelemetry span for tracing
budget_id: Budget row to load instead of the proxy-wide default
Returns:
LiteLLM_BudgetTable if configured and found, None otherwise
"""
if prisma_client is None or litellm.max_end_user_budget_id is None:
default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id
if prisma_client is None or default_budget_id is None:
return None
cache_key: Final = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
cache_key: Final = f"default_end_user_budget:{default_budget_id}"
# Check cache first
cached_budget: Final = await user_api_key_cache.async_get_cache(
@ -1388,12 +1403,13 @@ async def get_default_end_user_budget(
# Fetch from database
try:
budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique(
where={"budget_id": litellm.max_end_user_budget_id}
where={"budget_id": default_budget_id} # mutable-ok: prisma where clause
)
if budget_record is None:
verbose_proxy_logger.warning(
"Default end user budget not found in database: %s", litellm.max_end_user_budget_id
"Default end user budget not found in database: %s",
default_budget_id.replace("\r", "").replace("\n", ""),
)
return None
@ -1469,47 +1485,81 @@ async def get_team_member_default_budget(
return budget
async def resolve_default_end_user_budget(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
key_end_user_budget_id: str | None,
parent_otel_span: Span | None = None,
) -> LiteLLM_BudgetTable | None:
"""
The default budget for an end user with no budget of its own.
The key's ``end_user_budget_id`` takes precedence over the proxy-wide
``litellm.max_end_user_budget_id``; the proxy-wide default is the fallback when the key
names no budget or its budget row is missing.
"""
if key_end_user_budget_id is not None:
key_budget: Final = await get_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
budget_id=key_end_user_budget_id,
)
if key_budget is not None:
return key_budget
if litellm.max_end_user_budget_id is None:
return None
return await get_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
async def _apply_default_budget_to_end_user(
end_user_obj: LiteLLM_EndUserTable,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None = None,
key_end_user_budget_id: str | None = None,
) -> LiteLLM_EndUserTable:
"""
Helper function to apply default budget to end user if they don't have a budget assigned.
Returns the end user with the resolved default budget when it has no budget of its own.
A row whose own ``budget_id`` resolved to a budget is returned unchanged. Otherwise the
default is resolved on every call and set on a copy: the cached row carries at most the
proxy-wide default (readers such as the Prometheus customer gauges rely on that), never a
key's, so requests through keys with different defaults never observe each other's budget.
Args:
end_user_obj: The end user object to potentially apply default budget to
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving data
parent_otel_span: Optional OpenTelemetry span for tracing
Returns:
Updated end user object with default budget applied if applicable
key_end_user_budget_id: The requesting key's ``end_user_budget_id``, if any
"""
# If end user already has a budget assigned, no need to apply default
if end_user_obj.litellm_budget_table is not None:
if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None:
return end_user_obj
# If no default budget configured, return as-is
if litellm.max_end_user_budget_id is None:
if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None:
return end_user_obj
# Fetch and apply default budget
default_budget: Final = await get_default_end_user_budget(
default_budget: Final = await resolve_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
key_end_user_budget_id=key_end_user_budget_id,
parent_otel_span=parent_otel_span,
)
if default_budget is not None:
# Apply default budget to end user object
end_user_obj.litellm_budget_table = default_budget
verbose_proxy_logger.debug(
"Applied default budget %s to end user %s", litellm.max_end_user_budget_id, end_user_obj.user_id
)
if default_budget is None:
return end_user_obj
return end_user_obj
verbose_proxy_logger.debug(
"Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id
)
return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget}))
async def _check_end_user_budget(
@ -1714,6 +1764,7 @@ async def _end_user_is_known_unrestricted(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
token_end_user_max_budget: float | None,
key_end_user_budget_id: str | None = None,
) -> bool:
"""
True when the cached registry proves the id restricts nothing, so its row need not be read.
@ -1721,13 +1772,14 @@ async def _end_user_is_known_unrestricted(
Every field ``get_end_user_object`` callers consume (budget, spend under that budget, region,
default model, object permission, blocked) is part of the registry predicate, so an id outside
it is indistinguishable from one with no row at all. The skip is off whenever mere existence of
the row is meaningful: ``max_end_user_budget_id`` grafts a default budget onto any row that
exists, ``validate_end_user_id_in_db`` rejects ids that resolve to no row, and a token-supplied
``end_user_max_budget`` (a ``user_custom_auth`` callable can set one against an otherwise
unrestricted row) is enforced against the row's recorded spend.
the row is meaningful: ``max_end_user_budget_id`` or the key's ``end_user_budget_id`` grafts a
default budget onto any row that exists, ``validate_end_user_id_in_db`` rejects ids that resolve
to no row, and a token-supplied ``end_user_max_budget`` (a ``user_custom_auth`` callable can set
one against an otherwise unrestricted row) is enforced against the row's recorded spend.
"""
if (
litellm.max_end_user_budget_id is not None
or key_end_user_budget_id is not None
or litellm.validate_end_user_id_in_db
or token_end_user_max_budget is not None
):
@ -1749,12 +1801,13 @@ async def get_end_user_object(
parent_otel_span: Span | None = None,
proxy_logging_obj: ProxyLogging | None = None,
token_end_user_max_budget: float | None = None,
key_end_user_budget_id: str | None = None,
) -> LiteLLM_EndUserTable | None:
"""
Returns end user object from database or cache.
If end user exists but has no budget_id, applies the default budget
(if configured via litellm.max_end_user_budget_id).
If end user exists but has no budget_id, applies the default budget: the key's
``end_user_budget_id`` when set, otherwise ``litellm.max_end_user_budget_id``.
Args:
end_user_id: The ID of the end user
@ -1766,6 +1819,7 @@ async def get_end_user_object(
token_end_user_max_budget: ``valid_token.end_user_max_budget``, when the caller holds a
token. Budget enforcement reads the row's spend, so a row that restricts nothing on
its own must still be loaded when the token carries a budget for it.
key_end_user_budget_id: The requesting key's default end-user budget, if any
Returns:
LiteLLM_EndUserTable if found, None otherwise
@ -1784,22 +1838,20 @@ async def get_end_user_object(
model_type=LiteLLM_EndUserTable,
)
if cached_user_obj is not None:
return_obj = cached_user_obj
# Apply default budget if needed
return_obj = await _apply_default_budget_to_end_user(
end_user_obj=return_obj,
return await _apply_default_budget_to_end_user(
end_user_obj=cached_user_obj,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
key_end_user_budget_id=key_end_user_budget_id,
)
return return_obj
if await _end_user_is_known_unrestricted(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
token_end_user_max_budget=token_end_user_max_budget,
key_end_user_budget_id=key_end_user_budget_id,
):
return None
@ -1813,26 +1865,30 @@ async def get_end_user_object(
if response is None:
raise Exception
# Convert to LiteLLM_EndUserTable object
_response = LiteLLM_EndUserTable.model_validate(response.dict())
# Apply default budget if needed
_response = await _apply_default_budget_to_end_user(
end_user_obj=_response,
end_user_row: Final = await _apply_default_budget_to_end_user(
end_user_obj=LiteLLM_EndUserTable.model_validate(response.dict()),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
# Save to cache
await user_api_key_cache.async_set_cache(
key=_key,
value=_response,
value=end_user_row,
model_type=LiteLLM_EndUserTable,
ttl=get_management_object_ttl(user_api_key_cache),
)
return _response
if key_end_user_budget_id is None:
return end_user_row
return await _apply_default_budget_to_end_user(
end_user_obj=end_user_row,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
key_end_user_budget_id=key_end_user_budget_id,
)
except Exception:
return None
@ -1849,6 +1905,7 @@ async def resolve_and_validate_end_user_id(
parent_otel_span: Span | None = None,
proxy_logging_obj: ProxyLogging | None = None,
route: str = "",
key_end_user_budget_id: str | None = None,
) -> str | None:
"""Optionally drop end-user ids that don't resolve to a known DB row.
@ -1862,9 +1919,10 @@ async def resolve_and_validate_end_user_id(
- LiteLLM_UserTable.user_id
- LiteLLM_UserTable.user_email (case-insensitive)
If the id doesn't match but ``litellm.max_end_user_budget_id`` is set,
we still preserve the id so the default end-user budget is applied
downstream; otherwise we return None.
If the id doesn't match but a default end-user budget is configured
(``litellm.max_end_user_budget_id`` or the key's ``end_user_budget_id``),
we still preserve the id so that budget is applied downstream; otherwise
we return None.
DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they
share the same cache as the rest of the auth path instead of adding new
@ -1877,12 +1935,13 @@ async def resolve_and_validate_end_user_id(
if prisma_client is None:
return raw_end_user_id
has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None
cache_key: Final = f"end_user_validation:{raw_end_user_id}"
cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key)
if cached == "valid":
return raw_end_user_id
if cached == "invalid":
return raw_end_user_id if litellm.max_end_user_budget_id else None
return raw_end_user_id if has_default_budget else None
is_valid: Final = await _end_user_id_exists_in_db(
end_user_id=raw_end_user_id,
@ -1899,12 +1958,7 @@ async def resolve_and_validate_end_user_id(
ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL),
)
if is_valid:
return raw_end_user_id
# Preserve id so the caller can still apply litellm.max_end_user_budget_id.
if litellm.max_end_user_budget_id:
return raw_end_user_id
return None
return raw_end_user_id if is_valid or has_default_budget else None
async def _end_user_id_exists_in_db(

View file

@ -56,6 +56,7 @@ from litellm.proxy.auth.auth_checks import (
common_checks,
get_end_user_object,
get_jwt_key_mapping_object,
get_key_end_user_budget_id,
get_object_permission,
get_project_object,
get_team_membership,
@ -64,6 +65,7 @@ from litellm.proxy.auth.auth_checks import (
is_valid_fallback_model,
jwt_key_mapping_cache_key,
resolve_and_validate_end_user_id,
resolve_default_end_user_budget,
)
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
from litellm.proxy.auth.auth_method import AuthMethod
@ -2680,6 +2682,7 @@ async def _run_centralized_common_checks(
# resolved the end-user id and attached it here. Reuse that to avoid a
# second extraction pass; fall back to extracting locally when the
# function is invoked in isolation (e.g. in direct unit tests).
key_end_user_budget_id: Final = get_key_end_user_budget_id(user_api_key_auth_obj.metadata)
end_user_id = user_api_key_auth_obj.end_user_id
if end_user_id is None:
raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
@ -2690,7 +2693,10 @@ async def _run_centralized_common_checks(
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
key_end_user_budget_id=key_end_user_budget_id,
)
if end_user_id is not None and key_end_user_budget_id is not None:
user_api_key_auth_obj.end_user_id = end_user_id
fetch_coros: Final = []
if user_api_key_auth_obj.team_id is not None and user_api_key_auth_obj.team_id != UI_TEAM_ID:
@ -2753,6 +2759,7 @@ async def _run_centralized_common_checks(
proxy_logging_obj=proxy_logging_obj,
route=route,
token_end_user_max_budget=user_api_key_auth_obj.end_user_max_budget,
key_end_user_budget_id=key_end_user_budget_id,
),
)
)
@ -2857,6 +2864,17 @@ async def _run_centralized_common_checks(
user_api_key_auth_obj.project_metadata = project_object.metadata
user_api_key_auth_obj.project_alias = project_object.project_alias
if end_user_id and key_end_user_budget_id is not None and prisma_client is not None:
await _apply_key_end_user_default_budget_to_token(
valid_token=user_api_key_auth_obj,
end_user_object=end_user_object,
key_end_user_budget_id=key_end_user_budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
keep_token_limits=user_custom_auth is not None,
)
skip_budget_checks: Final = _should_skip_budget_checks(
request_data=request_data,
route=route,
@ -2945,6 +2963,46 @@ async def _noop_none() -> None:
return
async def _apply_key_end_user_default_budget_to_token(
valid_token: UserAPIKeyAuth,
end_user_object: LiteLLM_EndUserTable | None,
key_end_user_budget_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
keep_token_limits: bool,
) -> None:
"""The builder's end-user pass runs before the key is resolved, so only here can the key's
``end_user_budget_id`` win over the proxy-wide default on the token that reservation reads.
On the virtual-key path the token's end-user limits are the builder's proxy-wide defaults and
the key budget replaces them wholesale. With ``keep_token_limits`` (custom auth) the token's
limits are caps the custom auth callable set, so the key budget only fills the ones it left
unset."""
default_budget: Final = (
end_user_object.litellm_budget_table
if end_user_object is not None
else await resolve_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
key_end_user_budget_id=key_end_user_budget_id,
parent_otel_span=parent_otel_span,
)
)
if default_budget is None:
return
if not keep_token_limits or valid_token.end_user_max_budget is None:
valid_token.end_user_max_budget = default_budget.max_budget
if not keep_token_limits or valid_token.end_user_tpm_limit is None:
valid_token.end_user_tpm_limit = default_budget.tpm_limit
if not keep_token_limits or valid_token.end_user_rpm_limit is None:
valid_token.end_user_rpm_limit = default_budget.rpm_limit
if not keep_token_limits or valid_token.end_user_tpd_limit is None:
valid_token.end_user_tpd_limit = default_budget.tpd_limit
if not keep_token_limits or valid_token.end_user_model_max_budget is None:
valid_token.end_user_model_max_budget = default_budget.model_max_budget
async def _reserve_budget_after_common_checks(
user_api_key_auth_obj: UserAPIKeyAuth,
request_data: dict,
@ -3094,6 +3152,7 @@ async def _authorize_authenticated_request(
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
key_end_user_budget_id=get_key_end_user_budget_id(user_api_key_auth_obj.metadata),
)
if resolved_end_user_id is not None:
user_api_key_auth_obj.end_user_id = resolved_end_user_id
@ -3371,6 +3430,7 @@ async def _lookup_end_user_and_apply_budget(
):
"""Look up end_user from DB and apply budget limits to valid_token."""
end_user_object = None
key_end_user_budget_id: Final = get_key_end_user_budget_id(valid_token.metadata)
try:
end_user_object = await get_end_user_object(
end_user_id=valid_token.end_user_id,
@ -3380,6 +3440,7 @@ async def _lookup_end_user_and_apply_budget(
proxy_logging_obj=proxy_logging_obj,
route=route,
token_end_user_max_budget=valid_token.end_user_max_budget,
key_end_user_budget_id=key_end_user_budget_id,
)
if end_user_object is not None:
end_user_params = {
@ -3395,12 +3456,11 @@ async def _lookup_end_user_and_apply_budget(
valid_token = update_valid_token_with_end_user_params(
valid_token=valid_token, end_user_params=end_user_params
)
elif litellm.max_end_user_budget_id is not None:
from litellm.proxy.auth.auth_checks import get_default_end_user_budget
default_budget: Final = await get_default_end_user_budget(
elif key_end_user_budget_id is not None or litellm.max_end_user_budget_id is not None:
default_budget: Final = await resolve_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
key_end_user_budget_id=key_end_user_budget_id,
parent_otel_span=parent_otel_span,
)
if default_budget is not None:
@ -3413,6 +3473,8 @@ async def _lookup_end_user_and_apply_budget(
valid_token = update_valid_token_with_end_user_params(
valid_token=valid_token, end_user_params=end_user_params
)
if valid_token.end_user_max_budget is None:
valid_token.end_user_max_budget = default_budget.max_budget
except Exception as e:
if isinstance(e, litellm.BudgetExceededError):
raise e

View file

@ -55,6 +55,7 @@ from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
can_team_access_model,
get_jwt_key_mapping_cache_keys_for_token,
get_key_end_user_budget_id,
get_org_object,
get_project_object,
get_team_object,
@ -1175,6 +1176,13 @@ async def _common_key_generation_helper(
detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."},
)
await _validate_end_user_budget_id_change(
requested_budget_id=_requested_end_user_budget_id(data),
existing_budget_id=None,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
enforce_output_token_estimates_are_admin_only(
data=data,
existing_metadata=None,
@ -1930,6 +1938,7 @@ async def generate_key_fn(
- organization_id: Optional[str] - The organization id of the key. If not set, and team_id is set, the organization id will be the same as the team id. If conflict, an error will be raised.
- project_id: Optional[str] - The project id of the key. When set, models and max_budget are validated against the project's limits.
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
- end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Takes precedence over `litellm_settings.max_end_user_budget_id`.
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
- config: Optional[dict] - any key-specific configs, overrides config in config.yaml
@ -2142,6 +2151,7 @@ async def generate_service_account_key_fn(
- team_id: Optional[str] - The team id of the key
- user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
- end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it.
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
- config: Optional[dict] - any key-specific configs, overrides config in config.yaml
@ -2887,6 +2897,40 @@ def _require_prisma_client(prisma_client: PrismaClient | None) -> PrismaClient:
return prisma_client
def _requested_end_user_budget_id(data: KeyRequestBase) -> str | None:
"""A ``metadata`` body replaces the stored metadata wholesale, so one without the field clears it."""
if data.end_user_budget_id is not None:
return data.end_user_budget_id
if data.metadata is None:
return None
return get_key_end_user_budget_id(data.metadata) or ""
async def _validate_end_user_budget_id_change(
requested_budget_id: str | None,
existing_budget_id: str | None,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
) -> None:
"""A key's default end-user budget overrides the proxy-wide one, so only proxy admins
may change it, and a non-empty value must name an existing budget (empty clears it)."""
if requested_budget_id is None or requested_budget_id == (existing_budget_id or ""):
return
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
forbidden_detail: Final = { # mutable-ok: FastAPI detail contract
"error": "Only proxy admins can set end_user_budget_id on a key."
}
raise HTTPException(status_code=403, detail=forbidden_detail)
if requested_budget_id == "":
return
budget_row: Final = await BudgetRepository(_require_prisma_client(prisma_client)).find_by_id(requested_budget_id)
if budget_row is None:
missing_detail: Final = { # mutable-ok: FastAPI detail contract
"error": f"end_user_budget_id={requested_budget_id} does not match any budget."
}
raise HTTPException(status_code=400, detail=missing_detail)
async def _validate_update_key_data(
data: UpdateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
@ -2995,6 +3039,15 @@ async def _validate_update_key_data(
detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."},
)
await _validate_end_user_budget_id_change(
requested_budget_id=_requested_end_user_budget_id(data),
existing_budget_id=get_key_end_user_budget_id(
_existing_metadata if isinstance(_existing_metadata, dict) else None
),
user_api_key_dict=user_api_key_dict,
prisma_client=checked_prisma_client,
)
enforce_output_token_estimates_are_admin_only(
data=data,
existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None,
@ -3182,6 +3235,7 @@ async def update_key_fn(
- project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected.
- organization_id: Optional[str] - The organization id of the key.
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
- end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it.
- models: Optional[list] - Model_name's a user is allowed to call
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
@ -5383,6 +5437,14 @@ async def _execute_virtual_key_regeneration(
user_api_key_dict=user_api_key_dict,
entity="key",
)
await _validate_end_user_budget_id_change(
requested_budget_id=_requested_end_user_budget_id(data),
existing_budget_id=get_key_end_user_budget_id(
_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None
),
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
new_token: Final = await get_new_token(data=data)
new_token_hash: Final = hash_token(new_token)

View file

@ -1,5 +1,6 @@
import asyncio
import json
from collections.abc import Mapping
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal, Optional
from unittest.mock import AsyncMock, MagicMock, patch
@ -4779,6 +4780,28 @@ async def test_resolve_end_user_preserves_id_when_default_budget_configured(_val
assert result == "new-customer"
@pytest.mark.asyncio
@pytest.mark.parametrize("cached_verdict", [None, "invalid"])
async def test_resolve_end_user_preserves_id_when_only_the_key_default_budget_is_configured(
_validate_flag_on, monkeypatch, cached_verdict
):
"""With no proxy-wide default, a key-level end_user_budget_id still keeps an unregistered id
alive so the key's budget can be applied to that new customer downstream."""
from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id
_patch_validation_helpers(monkeypatch)
cache = _validation_cache()
cache.async_get_cache = AsyncMock(return_value=cached_verdict)
result = await resolve_and_validate_end_user_id(
raw_end_user_id="new-customer",
prisma_client=MagicMock(),
user_api_key_cache=cache,
key_end_user_budget_id="svc-a-budget",
)
assert result == "new-customer"
@pytest.mark.asyncio
async def test_resolve_end_user_drops_unknown_email(_validate_flag_on, monkeypatch):
from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id
@ -6608,6 +6631,208 @@ async def test_get_end_user_object_token_budget_gate_keeps_fetching_unrestricted
mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited()
def _budget_lookup_by_id(budgets: Mapping[str, float]) -> AsyncMock:
"""A ``litellm_budgettable.find_unique`` double that serves the given budgets by id."""
async def _find_unique(where: Mapping[str, str]) -> MagicMock | None:
budget_id = where["budget_id"]
if budget_id not in budgets:
return None
row = MagicMock()
row.dict = lambda: {"budget_id": budget_id, "max_budget": budgets[budget_id]}
return row
return AsyncMock(side_effect=_find_unique)
@pytest.mark.asyncio
async def test_get_end_user_object_key_default_budget_beats_global_default_without_leaking_across_keys(
monkeypatch,
):
"""Two service-account keys with different ``end_user_budget_id`` values must each see their
own default on the same unknown-but-existing end user, and the proxy-wide default must lose
to both. The row is cached after the first call, so the second call exercises the cache path.
"""
from litellm.proxy.auth.auth_checks import get_end_user_object
monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-eu-budget")
monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False)
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-shared"))
mock_prisma.db.litellm_budgettable.find_unique = _budget_lookup_by_id(
{"global-eu-budget": 100.0, "svc-a-budget": 0.5, "svc-b-budget": 7.0}
)
cache = UserApiKeyCache()
for_key_a = await get_end_user_object(
end_user_id="eu-shared",
prisma_client=mock_prisma,
user_api_key_cache=cache,
key_end_user_budget_id="svc-a-budget",
)
for_key_b = await get_end_user_object(
end_user_id="eu-shared",
prisma_client=mock_prisma,
user_api_key_cache=cache,
key_end_user_budget_id="svc-b-budget",
)
for_plain_key = await get_end_user_object(
end_user_id="eu-shared",
prisma_client=mock_prisma,
user_api_key_cache=cache,
)
assert for_key_a is not None and for_key_a.litellm_budget_table is not None
assert for_key_a.litellm_budget_table.max_budget == 0.5
assert for_key_b is not None and for_key_b.litellm_budget_table is not None
assert for_key_b.litellm_budget_table.max_budget == 7.0
assert for_plain_key is not None and for_plain_key.litellm_budget_table is not None
assert for_plain_key.litellm_budget_table.max_budget == 100.0
mock_prisma.db.litellm_endusertable.find_unique.assert_awaited_once()
@pytest.mark.asyncio
async def test_get_end_user_object_cached_row_does_not_carry_another_keys_default_budget(monkeypatch):
"""A key without a default must see the end user unrestricted even after a key with a default
populated the shared per-end-user cache entry for the same id."""
from litellm.proxy.auth.auth_checks import get_end_user_object
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False)
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-shared"))
mock_prisma.db.litellm_budgettable.find_unique = _budget_lookup_by_id({"svc-a-budget": 0.5})
cache = UserApiKeyCache()
for_key_a = await get_end_user_object(
end_user_id="eu-shared",
prisma_client=mock_prisma,
user_api_key_cache=cache,
key_end_user_budget_id="svc-a-budget",
)
for_plain_key = await get_end_user_object(
end_user_id="eu-shared",
prisma_client=mock_prisma,
user_api_key_cache=cache,
)
assert for_key_a is not None and for_key_a.litellm_budget_table is not None
assert for_key_a.litellm_budget_table.max_budget == 0.5
assert for_plain_key is not None
assert for_plain_key.litellm_budget_table is None
@pytest.mark.asyncio
async def test_get_end_user_object_caches_row_with_global_default_but_never_a_key_default(monkeypatch):
"""The cached row is what post-request readers (Prometheus customer gauges) see: it must keep
the proxy-wide default exactly as before, while a key default stays on the request copy."""
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key
monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-budget")
monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False)
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-cached"))
mock_prisma.db.litellm_budgettable.find_unique = _budget_lookup_by_id({"svc-a-budget": 0.5, "global-budget": 7.0})
cache = UserApiKeyCache()
for_key_a = await get_end_user_object(
end_user_id="eu-cached",
prisma_client=mock_prisma,
user_api_key_cache=cache,
key_end_user_budget_id="svc-a-budget",
)
cached = await cache.async_get_cache(key=end_user_cache_key("eu-cached"), model_type=LiteLLM_EndUserTable)
assert for_key_a is not None and for_key_a.litellm_budget_table is not None
assert for_key_a.litellm_budget_table.max_budget == 0.5
assert cached is not None and cached.litellm_budget_table is not None
assert cached.litellm_budget_table.max_budget == 7.0
@pytest.mark.asyncio
async def test_get_end_user_object_key_default_budget_loads_unrestricted_row_without_global_default(
end_user_registry_skip_enabled,
):
"""With no proxy-wide default, a key default alone must keep the registry skip off, otherwise
the unrestricted row is never loaded and the key default is never enforced.
"""
from litellm.proxy.auth.auth_checks import get_end_user_object
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1", spend=3.0))
mock_prisma.db.litellm_budgettable.find_unique = _budget_lookup_by_id({"svc-a-budget": 2.0})
result = await get_end_user_object(
end_user_id="eu-anon-1",
prisma_client=mock_prisma,
user_api_key_cache=UserApiKeyCache(),
key_end_user_budget_id="svc-a-budget",
)
assert result is not None
assert result.spend == 3.0
assert result.litellm_budget_table is not None
assert result.litellm_budget_table.max_budget == 2.0
mock_prisma.db.litellm_endusertable.find_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_get_end_user_object_explicit_end_user_budget_beats_key_default(monkeypatch):
from litellm.proxy.auth.auth_checks import get_end_user_object
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False)
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(
return_value=_end_user_db_row(
"eu-vip",
budget_id="vip-budget",
litellm_budget_table={"budget_id": "vip-budget", "max_budget": 500.0},
)
)
mock_prisma.db.litellm_budgettable.find_unique = _budget_lookup_by_id({"svc-a-budget": 0.5})
result = await get_end_user_object(
end_user_id="eu-vip",
prisma_client=mock_prisma,
user_api_key_cache=UserApiKeyCache(),
key_end_user_budget_id="svc-a-budget",
)
assert result is not None and result.litellm_budget_table is not None
assert result.litellm_budget_table.max_budget == 500.0
mock_prisma.db.litellm_budgettable.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_default_end_user_budget_falls_back_to_global_when_key_budget_is_missing(monkeypatch):
from litellm.proxy.auth.auth_checks import resolve_default_end_user_budget
monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-eu-budget")
mock_prisma = MagicMock()
mock_prisma.db.litellm_budgettable.find_unique = _budget_lookup_by_id({"global-eu-budget": 100.0})
resolved = await resolve_default_end_user_budget(
prisma_client=mock_prisma,
user_api_key_cache=UserApiKeyCache(),
key_end_user_budget_id="deleted-budget",
)
assert resolved is not None
assert resolved.budget_id == "global-eu-budget"
assert resolved.max_budget == 100.0
@pytest.mark.asyncio
async def test_end_user_id_validation_gate_still_resolves_unrestricted_end_users(monkeypatch):
"""

View file

@ -222,6 +222,117 @@ async def test_custom_auth_token_budget_still_loads_and_caches_unrestricted_end_
assert await cache.async_get_cache(key=end_user_cache_key("customer-1")) is not None
@pytest.mark.asyncio
async def test_custom_auth_key_default_end_user_budget_reaches_the_token_for_a_new_end_user(monkeypatch):
"""A custom-auth token that carries a key ``end_user_budget_id`` must enforce that budget on a
brand-new end user, ahead of the proxy-wide default, from the very first request."""
from unittest.mock import MagicMock
from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-eu-budget")
budgets = {"global-eu-budget": 100.0, "svc-a-budget": 0.5}
async def _find_budget(where):
row = MagicMock()
row.dict = lambda: {"budget_id": where["budget_id"], "max_budget": budgets[where["budget_id"]]}
return row
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(side_effect=_find_budget)
valid_token, end_user_object = await _lookup_end_user_and_apply_budget(
valid_token=UserAPIKeyAuth(
token="test_token",
end_user_id="customer-new",
metadata={"end_user_budget_id": "svc-a-budget"},
),
route="/v1/chat/completions",
parent_otel_span=None,
prisma_client=mock_prisma,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=MagicMock(),
)
assert end_user_object is None
assert valid_token.end_user_max_budget == 0.5
@pytest.mark.asyncio
async def test_custom_auth_cap_stays_below_the_key_default_end_user_budget(monkeypatch):
"""A custom auth callable that already capped the end user tighter than the key's default
budget keeps its cap: the key default never loosens what custom auth set."""
from unittest.mock import MagicMock
from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
async def _find_budget(where):
row = MagicMock()
row.dict = lambda: {"budget_id": where["budget_id"], "max_budget": 0.5}
return row
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(side_effect=_find_budget)
valid_token, _ = await _lookup_end_user_and_apply_budget(
valid_token=UserAPIKeyAuth(
token="test_token",
end_user_id="customer-new",
end_user_max_budget=0.1,
metadata={"end_user_budget_id": "svc-a-budget"},
),
route="/v1/chat/completions",
parent_otel_span=None,
prisma_client=mock_prisma,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=MagicMock(),
)
assert valid_token.end_user_max_budget == 0.1
@pytest.mark.asyncio
async def test_custom_auth_proxy_wide_default_end_user_budget_reaches_an_uncapped_token(monkeypatch):
"""With no key default, a brand-new end user on a custom-auth token that set no cap gets the
proxy-wide default budget's cap, the same way the virtual-key path already applies it."""
from unittest.mock import MagicMock
from litellm.proxy.auth.user_api_key_auth import _lookup_end_user_and_apply_budget
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-eu-budget")
async def _find_budget(where):
row = MagicMock()
row.dict = lambda: {"budget_id": where["budget_id"], "max_budget": 100.0}
return row
mock_prisma = MagicMock()
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(side_effect=_find_budget)
valid_token, end_user_object = await _lookup_end_user_and_apply_budget(
valid_token=UserAPIKeyAuth(token="test_token", end_user_id="customer-new"),
route="/v1/chat/completions",
parent_otel_span=None,
prisma_client=mock_prisma,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=MagicMock(),
)
assert end_user_object is None
assert valid_token.end_user_max_budget == 100.0
def test_update_valid_token_does_not_override_custom_auth_values_with_none():
"""
Greptile feedback: if custom auth sets end_user_model_max_budget on the token,

View file

@ -4,6 +4,7 @@ import logging
import os
import subprocess
import sys
from collections.abc import Mapping
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from functools import partial
@ -4374,6 +4375,186 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t
}
def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock:
row = MagicMock()
row.dict = lambda: {"budget_id": budget_id, "max_budget": max_budget}
return row
async def _run_centralized_checks_with_key_end_user_budget(
token: UserAPIKeyAuth,
end_user_row: MagicMock | None,
budgets: Mapping[str, float],
request_user: str | None = None,
user_api_key_cache: DualCache | None = None,
custom_auth: bool = False,
) -> UserAPIKeyAuth:
"""Run the centralized checks with a fake DB and return the token handed to budget reservation.
With ``custom_auth`` the token stands for one a custom auth callable returned and the checks
run under ``custom_auth_run_common_checks``."""
from fastapi import Request
from starlette.datastructures import URL
import litellm.proxy.proxy_server as _proxy_server_mod
async def _find_budget(where: Mapping[str, str]) -> MagicMock | None:
budget_id = where["budget_id"]
return _end_user_budget_row(budget_id, budgets[budget_id]) if budget_id in budgets else None
prisma_client = MagicMock()
prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_endusertable.find_unique = AsyncMock(return_value=end_user_row)
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(side_effect=_find_budget)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
attrs = {
**_proxy_attrs_for_centralized_checks(
user_custom_auth=AsyncMock() if custom_auth else None, flag=custom_auth
),
"prisma_client": prisma_client,
"user_api_key_cache": user_api_key_cache if user_api_key_cache is not None else DualCache(),
"proxy_logging_obj": proxy_logging_obj,
}
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
with (
patch( # test-quality-ok: the authz gate has its own tests above; this one checks what reaches reservation
"litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
),
patch( # test-quality-ok: reservation is the observable boundary; its input token is what is asserted
"litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks",
new_callable=AsyncMock,
) as mock_reserve,
):
await _run_centralized_common_checks(
user_api_key_auth_obj=token,
request=request,
request_data={"model": "gpt-5.4-mini", "user": request_user or token.end_user_id},
route="/chat/completions",
)
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
mock_reserve.assert_awaited_once()
return mock_reserve.call_args.kwargs["user_api_key_auth_obj"]
@pytest.mark.asyncio
async def test_centralized_common_checks_keeps_a_validated_away_end_user_when_the_key_has_a_default(monkeypatch):
"""With ``validate_end_user_id_in_db`` on and no proxy-wide default, the builder drops an
unregistered customer id before it knows the key. The central gate must re-resolve it with the
key's default so the customer is both budgeted and attributed on the first request."""
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True)
cache = DualCache()
await cache.async_set_cache(key="end_user_validation:cust-new", value="invalid")
token = UserAPIKeyAuth(
api_key="sk-test",
token="hashed",
end_user_id=None,
metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"},
)
reserved_token = await _run_centralized_checks_with_key_end_user_budget(
token, end_user_row=None, budgets={"svc-a-budget": 0.5}, request_user="cust-new", user_api_key_cache=cache
)
assert reserved_token.end_user_id == "cust-new"
assert reserved_token.end_user_max_budget == 0.5
@pytest.mark.asyncio
async def test_centralized_common_checks_reserves_key_default_budget_for_a_brand_new_end_user(monkeypatch):
"""A service-account key's ``end_user_budget_id`` must reach the token before the budget
reservation runs, on the very first request, when no end-user row exists yet and even though
the builder already applied the proxy-wide default."""
monkeypatch.setattr(litellm, "max_end_user_budget_id", "global-eu-budget")
token = UserAPIKeyAuth(
api_key="sk-test",
token="hashed",
end_user_id="cust-new",
end_user_max_budget=100.0,
metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"},
)
reserved_token = await _run_centralized_checks_with_key_end_user_budget(
token, end_user_row=None, budgets={"global-eu-budget": 100.0, "svc-a-budget": 0.5}
)
assert reserved_token.end_user_max_budget == 0.5
@pytest.mark.asyncio
async def test_centralized_common_checks_keeps_an_end_users_own_budget_over_the_key_default(monkeypatch):
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
token = UserAPIKeyAuth(
api_key="sk-test",
token="hashed",
end_user_id="cust-vip",
end_user_max_budget=500.0,
metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"},
)
end_user_row = MagicMock()
end_user_row.dict = lambda: {
"user_id": "cust-vip",
"blocked": False,
"spend": 0.0,
"budget_id": "vip-budget",
"litellm_budget_table": {"budget_id": "vip-budget", "max_budget": 500.0},
}
reserved_token = await _run_centralized_checks_with_key_end_user_budget(
token, end_user_row=end_user_row, budgets={"svc-a-budget": 0.5}
)
assert reserved_token.end_user_max_budget == 500.0
@pytest.mark.asyncio
async def test_centralized_common_checks_keeps_a_stricter_custom_auth_cap_over_the_key_default(monkeypatch):
"""A custom auth callable that caps the end user tighter than the key's default budget keeps
its cap and its rate limit. The key default only fills the limits the callable left unset."""
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
token = UserAPIKeyAuth(
api_key="sk-test",
token="hashed",
end_user_id="cust-new",
end_user_max_budget=0.1,
end_user_rpm_limit=3,
metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"},
)
reserved_token = await _run_centralized_checks_with_key_end_user_budget(
token, end_user_row=None, budgets={"svc-a-budget": 0.5}, custom_auth=True
)
assert reserved_token.end_user_max_budget == 0.1
assert reserved_token.end_user_rpm_limit == 3
@pytest.mark.asyncio
async def test_centralized_common_checks_fills_a_custom_auth_token_without_a_cap_from_the_key_default(monkeypatch):
monkeypatch.setattr(litellm, "max_end_user_budget_id", None)
token = UserAPIKeyAuth(
api_key="sk-test",
token="hashed",
end_user_id="cust-new",
metadata={"service_account_id": "svc-a", "end_user_budget_id": "svc-a-budget"},
)
reserved_token = await _run_centralized_checks_with_key_end_user_budget(
token, end_user_row=None, budgets={"svc-a-budget": 0.5}, custom_auth=True
)
assert reserved_token.end_user_max_budget == 0.5
class _RecordingTeamModelBudgetLimiter:
def __init__(self):
self.calls = []

View file

@ -58,8 +58,10 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
_list_key_helper,
_persist_deleted_verification_tokens,
_process_single_key_update,
_requested_end_user_budget_id,
_save_deleted_verification_token_records,
_transform_verification_tokens_to_deleted_records,
_validate_end_user_budget_id_change,
_validate_max_budget,
_validate_reset_spend_value,
_validate_update_key_data,
@ -1869,6 +1871,202 @@ async def test_generate_key_throttle_allowed_for_admin():
assert mock_generate_key.called
@pytest.mark.asyncio
async def test_generate_key_end_user_budget_id_rejected_for_non_admin():
"""A key's default end-user budget overrides the proxy-wide one, so a non-admin must not
be able to pick a looser one for the customers their key creates."""
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock()
with pytest.raises(HTTPException) as exc:
await _validate_end_user_budget_id_change(
requested_budget_id="svc-a-budget",
existing_budget_id=None,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-alice",
user_id="alice",
),
prisma_client=mock_prisma_client,
)
assert int(getattr(exc.value, "status_code", 0)) == 403
assert "Only proxy admins can set end_user_budget_id" in str(exc.value.detail)
mock_prisma_client.db.litellm_budgettable.find_unique.assert_not_awaited()
await _validate_end_user_budget_id_change(
requested_budget_id="",
existing_budget_id=None,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-alice",
user_id="alice",
),
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_generate_key_end_user_budget_id_must_name_an_existing_budget():
"""A typo in end_user_budget_id would silently leave new customers on the proxy-wide default,
so key creation rejects an id that matches no budget row."""
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc:
await _validate_end_user_budget_id_change(
requested_budget_id="no-such-budget",
existing_budget_id=None,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
prisma_client=mock_prisma_client,
)
assert int(getattr(exc.value, "status_code", 0)) == 400
assert "no-such-budget" in str(exc.value.detail)
mock_prisma_client.db.litellm_budgettable.find_unique.assert_awaited_once_with(
where={"budget_id": "no-such-budget"}
)
@pytest.mark.asyncio
async def test_generate_key_end_user_budget_id_lands_in_key_metadata():
"""The typed end_user_budget_id field is stored in key metadata, which is where auth reads it."""
budget_row = MagicMock()
budget_row.model_dump.return_value = {"budget_id": "svc-a-budget", "max_budget": 0.5}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row)
with (
patch( # test-quality-ok: the helper reads proxy_server globals, no seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
),
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: read as a proxy_server global
patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: read as a proxy_server global
patch( # test-quality-ok: assertion is on the metadata handed to the db writer
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
) as mock_generate_key,
):
mock_generate_key.return_value = {
"key": "sk-test-key",
"expires": None,
"user_id": "admin",
"team_id": None,
}
await _common_key_generation_helper(
data=GenerateKeyRequest(end_user_budget_id="svc-a-budget"),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
litellm_changed_by=None,
team_table=None,
)
assert mock_generate_key.call_args.kwargs["metadata"] == {"end_user_budget_id": "svc-a-budget"}
@pytest.mark.asyncio
async def test_update_key_end_user_budget_id_folds_into_metadata_and_survives_omission():
"""/key/update with end_user_budget_id writes it into metadata; an update that omits the field
(the edit form only sends what changed) keeps the value the key already had."""
existing_key = LiteLLM_VerificationToken(token="hashed", metadata={"end_user_budget_id": "svc-a-budget"})
updated = await prepare_key_update_data(
data=UpdateKeyRequest(key="sk-1", end_user_budget_id="svc-b-budget"), existing_key_row=existing_key
)
assert updated["metadata"]["end_user_budget_id"] == "svc-b-budget"
untouched = await prepare_key_update_data(
data=UpdateKeyRequest(key="sk-1", key_alias="renamed"), existing_key_row=existing_key
)
assert untouched["metadata"]["end_user_budget_id"] == "svc-a-budget"
@pytest.mark.asyncio
async def test_update_key_clears_end_user_budget_id_with_empty_string():
"""Sending an empty end_user_budget_id detaches the key default without touching any budget row,
so auth falls back to the proxy-wide default for that key's customers."""
from litellm.proxy.auth.auth_checks import get_key_end_user_budget_id
existing_key = LiteLLM_VerificationToken(token="hashed", metadata={"end_user_budget_id": "svc-a-budget"})
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
await _validate_update_key_data(
data=UpdateKeyRequest(key="sk-1", end_user_budget_id=""),
existing_key_row=existing_key,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
llm_router=None,
premium_user=False,
prisma_client=mock_prisma_client,
user_api_key_cache=MagicMock(),
)
cleared = await prepare_key_update_data(
data=UpdateKeyRequest(key="sk-1", end_user_budget_id="", metadata={"end_user_budget_id": "svc-a-budget"}),
existing_key_row=existing_key,
)
mock_prisma_client.db.litellm_budgettable.find_unique.assert_not_awaited()
assert get_key_end_user_budget_id(cleared["metadata"]) is None
@pytest.mark.asyncio
async def test_update_key_metadata_body_without_end_user_budget_id_is_a_clear_for_non_admin():
"""/key/update replaces metadata wholesale, so a non-admin sending metadata that drops the field
would detach the key default; that must be refused like an explicit clear, while an admin may do it."""
existing_key = LiteLLM_VerificationToken(token="hashed", metadata={"end_user_budget_id": "svc-a-budget"})
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
non_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-alice", user_id="alice")
with pytest.raises(HTTPException) as exc:
await _validate_update_key_data(
data=UpdateKeyRequest(key="sk-1", metadata={"team": "ops"}),
existing_key_row=existing_key,
user_api_key_dict=non_admin,
llm_router=None,
premium_user=False,
prisma_client=mock_prisma_client,
user_api_key_cache=MagicMock(),
)
assert int(getattr(exc.value, "status_code", 0)) == 403
await _validate_end_user_budget_id_change(
requested_budget_id=_requested_end_user_budget_id(
UpdateKeyRequest(key="sk-1", metadata={"team": "ops", "end_user_budget_id": "svc-a-budget"})
),
existing_budget_id="svc-a-budget",
user_api_key_dict=non_admin,
prisma_client=mock_prisma_client,
)
await _validate_end_user_budget_id_change(
requested_budget_id=_requested_end_user_budget_id(UpdateKeyRequest(key="sk-1", metadata={"team": "ops"})),
existing_budget_id="svc-a-budget",
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
prisma_client=mock_prisma_client,
)
assert _requested_end_user_budget_id(UpdateKeyRequest(key="sk-1", key_alias="renamed")) is None
mock_prisma_client.db.litellm_budgettable.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_regenerate_key_end_user_budget_id_rejected_for_non_admin():
"""/key/regenerate also accepts key params, so a non-admin must not be able to use it to attach
a looser default customer budget that /key/generate and /key/update would refuse."""
from litellm.proxy._types import RegenerateKeyRequest
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock()
with pytest.raises(HTTPException) as exc:
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=LiteLLM_VerificationToken(token="hashed", user_id="alice"),
hashed_api_key="hashed",
key="hashed",
data=RegenerateKeyRequest(end_user_budget_id="svc-a-budget"),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-alice", user_id="alice"
),
litellm_changed_by=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
)
assert int(getattr(exc.value, "status_code", 0)) == 403
assert "Only proxy admins can set end_user_budget_id" in str(exc.value.detail)
mock_prisma_client.db.litellm_verificationtoken.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_service_account_requires_team_id():
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"})

View file

@ -0,0 +1,19 @@
"use client";
import { useQuery, type UseQueryResult } from "@tanstack/react-query";
import { apiClient } from "@/components/networking";
import { budgetKeys, type budgetItem } from "./useBudgets";
const BUDGET_OPTIONS_PATH = "/budget/list";
export const useBudgetOptions = (accessToken: string | null, enabled = true): UseQueryResult<budgetItem[]> => {
const queryOptions = {
queryKey: [...budgetKeys.all, "options"],
queryFn: () => apiClient.get<budgetItem[]>(BUDGET_OPTIONS_PATH, { accessToken }),
enabled: Boolean(accessToken) && enabled,
staleTime: 60_000,
};
return useQuery(queryOptions);
};

View file

@ -0,0 +1,61 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, expect, it, vi } from "vitest";
import { chooseSelectOption } from "../../../tests/test-utils";
import { EndUserBudgetSelect } from "./EndUserBudgetSelect";
const useBudgetOptions = vi.fn();
vi.mock("@/app/(dashboard)/hooks/budgets/useBudgetOptions", () => ({
useBudgetOptions: (...args: unknown[]) => useBudgetOptions(...args),
}));
const BUDGETS = [
{ budget_id: "svc-a-budget", max_budget: 0.5, budget_duration: "30d", created_at: "", updated_at: "" },
{ budget_id: "svc-b-budget", max_budget: null, budget_duration: null, created_at: "", updated_at: "" },
];
describe("EndUserBudgetSelect", () => {
it("lets an admin pick one of the proxy's budgets and reports its id", async () => {
useBudgetOptions.mockReturnValue({ data: BUDGETS });
const onChange = vi.fn();
const user = userEvent.setup();
render(<EndUserBudgetSelect accessToken="tok" value={null} onChange={onChange} canEdit />);
await chooseSelectOption(user, screen.getByRole("combobox", { name: "Default Customer Budget" }), /svc-a-budget/);
expect(onChange).toHaveBeenLastCalledWith("svc-a-budget");
expect(useBudgetOptions).toHaveBeenCalledWith("tok", true);
});
it("shows a budget's cap and reset window next to its id", async () => {
useBudgetOptions.mockReturnValue({ data: BUDGETS });
const user = userEvent.setup();
render(<EndUserBudgetSelect accessToken="tok" value={null} onChange={vi.fn()} canEdit />);
await user.click(screen.getByRole("combobox"));
expect(await screen.findByRole("option", { name: /svc-a-budget/ })).toHaveTextContent("$0.5, resets 30d");
});
it("clears to null so the edit form can send an explicit empty value", async () => {
useBudgetOptions.mockReturnValue({ data: BUDGETS });
const onChange = vi.fn();
const user = userEvent.setup();
render(<EndUserBudgetSelect accessToken="tok" value="svc-a-budget" onChange={onChange} canEdit />);
await user.click(screen.getByRole("button", { name: "Clear" }));
expect(onChange).toHaveBeenLastCalledWith(null);
});
it("keeps the stored budget visible but read-only for a user who cannot change it", () => {
useBudgetOptions.mockReturnValue({ data: undefined });
render(<EndUserBudgetSelect accessToken="tok" value="svc-a-budget" onChange={vi.fn()} canEdit={false} />);
const combobox = screen.getByRole("combobox", { name: "Default Customer Budget" });
expect(combobox).toHaveValue("svc-a-budget");
expect(combobox).toBeDisabled();
expect(useBudgetOptions).toHaveBeenLastCalledWith("tok", false);
});
});

View file

@ -0,0 +1,55 @@
"use client";
import React from "react";
import { useBudgetOptions } from "@/app/(dashboard)/hooks/budgets/useBudgetOptions";
import type { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets";
import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect";
export const END_USER_BUDGET_HINT =
"Reusable budget applied to every new customer (end user) this key creates via `user` or x-litellm-end-user-id. " +
"Overrides the proxy-wide max_end_user_budget_id; customers that already have their own budget keep it.";
interface EndUserBudgetSelectProps {
readonly id?: string;
readonly accessToken: string | null;
readonly value: string | null;
readonly onChange: (next: string | null) => void;
readonly canEdit: boolean;
}
const budgetSublabel = (budget: budgetItem): string | undefined => {
const parts = [
budget.max_budget != null ? `$${budget.max_budget}` : null,
budget.budget_duration ? `resets ${budget.budget_duration}` : null,
].filter((part): part is string => part !== null);
return parts.length > 0 ? parts.join(", ") : undefined;
};
export const EndUserBudgetSelect: React.FC<EndUserBudgetSelectProps> = ({
id,
accessToken,
value,
onChange,
canEdit,
}) => {
const { data: budgets } = useBudgetOptions(accessToken, canEdit);
const options: SearchSelectOption[] = (budgets ?? []).map((budget) => ({
label: budget.budget_id,
value: budget.budget_id,
sublabel: budgetSublabel(budget),
}));
return (
<SearchSelect
inputId={id}
aria-label="Default Customer Budget"
placeholder="No default budget"
emptyText="No budgets found. Create one under Budgets."
options={options}
value={value}
onValueChange={onChange}
disabled={!canEdit}
/>
);
};

View file

@ -0,0 +1,45 @@
import { describe, expect, it } from "vitest";
import { endUserBudgetIdUpdate, keyOffersEndUserBudget, storedEndUserBudgetId } from "./endUserBudgetPayload";
describe("keyOffersEndUserBudget", () => {
it("offers the control on service account keys and on keys that already carry a budget", () => {
expect(keyOffersEndUserBudget({ service_account_id: "svc-a" })).toBe(true);
expect(keyOffersEndUserBudget({ end_user_budget_id: "svc-a-budget" })).toBe(true);
});
it.each([undefined, null, {}, { service_account_id: "" }, { tags: ["x"] }])("hides it for %j", (metadata) => {
expect(keyOffersEndUserBudget(metadata)).toBe(false);
});
});
describe("storedEndUserBudgetId", () => {
it("reads the budget id a key applies to the customers it creates", () => {
expect(storedEndUserBudgetId({ service_account_id: "svc-a", end_user_budget_id: "svc-a-budget" })).toBe(
"svc-a-budget",
);
});
it.each([undefined, null, "not-an-object", [], {}, { end_user_budget_id: 7 }])(
"reads %j as no default budget",
(metadata) => {
expect(storedEndUserBudgetId(metadata)).toBe("");
},
);
});
describe("endUserBudgetIdUpdate", () => {
it("leaves the field off the payload when the selection matches the stored value", () => {
expect(endUserBudgetIdUpdate("svc-a-budget", "svc-a-budget")).toBeUndefined();
expect(endUserBudgetIdUpdate(null, "")).toBeUndefined();
});
it("sends the newly selected budget id", () => {
expect(endUserBudgetIdUpdate("svc-b-budget", "svc-a-budget")).toBe("svc-b-budget");
expect(endUserBudgetIdUpdate("svc-a-budget", "")).toBe("svc-a-budget");
});
it("sends an empty string so the backend clears a previously stored budget", () => {
expect(endUserBudgetIdUpdate(null, "svc-a-budget")).toBe("");
});
});

View file

@ -0,0 +1,15 @@
const metadataString = (metadata: unknown, key: string): string => {
if (metadata === null || typeof metadata !== "object" || Array.isArray(metadata)) return "";
const value = (metadata as Record<string, unknown>)[key];
return typeof value === "string" ? value : "";
};
export const storedEndUserBudgetId = (metadata: unknown): string => metadataString(metadata, "end_user_budget_id");
export const keyOffersEndUserBudget = (metadata: unknown): boolean =>
metadataString(metadata, "service_account_id") !== "" || storedEndUserBudgetId(metadata) !== "";
export const endUserBudgetIdUpdate = (selected: string | null, stored: string): string | undefined => {
const next = selected ?? "";
return next === stored ? undefined : next;
};

View file

@ -33,6 +33,11 @@ vi.mock("@/lib/toast", () => ({
},
}));
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => state.authorized }));
vi.mock("@/app/(dashboard)/hooks/budgets/useBudgetOptions", () => ({
useBudgetOptions: () => ({
data: [{ budget_id: "svc-a-budget", max_budget: 0.5, created_at: "", updated_at: "" }],
}),
}));
vi.mock("@/app/(dashboard)/hooks/useCan", () => ({
default: (capability: string) => state.can[capability] ?? true,
}));
@ -647,6 +652,46 @@ describe("CreateKey", () => {
expect(JSON.parse(String(payload.metadata))).toStrictEqual({ service_account_id: "svc-account-1" });
expect(payload).not.toHaveProperty("user_id");
});
it("sends the chosen default customer budget with a service account", async () => {
state.teams = [{ team_id: "team-1", team_alias: "Team One", models: [] }];
await openModal({ teams: state.teams as unknown as Team[] });
await userEvent.click(screen.getByRole("radio", { name: "Service Account" }));
await userEvent.type(await screen.findByLabelText(/Service Account ID/), "svc-account-1");
await userEvent.click(await screen.findByLabelText("Team"));
await userEvent.click(await screen.findByRole("option", { name: /Team One/ }));
await openSection(/Optional Settings/i);
await userEvent.click(await screen.findByRole("combobox", { name: "Default Customer Budget" }));
await userEvent.click(await screen.findByRole("option", { name: /svc-a-budget/ }));
await submit();
await waitFor(() => {
expect(vi.mocked(keyCreateServiceAccountCall)).toHaveBeenCalled();
});
const payload = vi.mocked(keyCreateServiceAccountCall).mock.calls[0][1] as Record<string, unknown>;
expect(payload).toHaveProperty("end_user_budget_id", "svc-a-budget");
});
it("offers the default customer budget only to admins creating a service account", async () => {
await openModal();
await openSection(/Optional Settings/i);
await screen.findByLabelText(/Max Budget/);
expect(screen.queryByRole("combobox", { name: "Default Customer Budget" })).not.toBeInTheDocument();
await userEvent.click(screen.getByRole("radio", { name: "Service Account" }));
expect(await screen.findByRole("combobox", { name: "Default Customer Budget" })).toBeInTheDocument();
});
it("hides the default customer budget from a non-admin creating a service account", async () => {
state.authorized = { ...state.authorized, userRole: "Internal User" };
await openModal();
await userEvent.click(screen.getByRole("radio", { name: "Service Account" }));
await openSection(/Optional Settings/i);
await screen.findByLabelText(/Max Budget/);
expect(screen.queryByRole("combobox", { name: "Default Customer Budget" })).not.toBeInTheDocument();
});
});
describe("required field validation", () => {

View file

@ -25,7 +25,7 @@ import { TagsInput } from "@/app/(dashboard)/guardrails/_components/content_filt
import { ChevronDown, Info } from "lucide-react";
import React, { useEffect, useMemo, useRef, useState } from "react";
import { type Control, useForm, useWatch, type UseFormSetValue } from "react-hook-form";
import { rolesWithWriteAccess } from "../../utils/roles";
import { isProxyAdminRole, rolesWithWriteAccess } from "../../utils/roles";
import AgentSelector from "../agent_management/AgentSelector";
import SkillSelector from "../skills/SkillSelector";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
@ -52,6 +52,7 @@ import OrganizationDropdown from "../common_components/OrganizationDropdown";
import ProjectDropdown from "../common_components/ProjectDropdown";
import { CreateUserButton } from "../CreateUserButton";
import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor";
import { END_USER_BUDGET_HINT, EndUserBudgetSelect } from "../key_team_helpers/EndUserBudgetSelect";
import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
import { ModelMaxBudget, ModelMaxBudgetEditor } from "../key_team_helpers/ModelMaxBudgetEditor";
import { TagRateLimitEditor, TagRateLimitEntry } from "../key_team_helpers/TagRateLimitEditor";
@ -1068,6 +1069,30 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
availableModels={modelsToPick}
/>
</Field>
{keyOwner === "service_account" && isProxyAdminRole(userRole ?? "") && (
<MountedFormField
className="mt-4"
label={
<span>
Default Customer Budget{" "}
<SimpleTooltip content={END_USER_BUDGET_HINT}>
<Info className="ml-1 inline size-3.5 align-text-bottom" />
</SimpleTooltip>
</span>
}
name="end_user_budget_id"
>
{(control) => (
<EndUserBudgetSelect
id={control.id}
accessToken={accessToken}
value={typeof control.value === "string" ? control.value : null}
onChange={control.onChange}
canEdit
/>
)}
</MountedFormField>
)}
<MountedFormField
className="mt-4"
label={

View file

@ -69,6 +69,15 @@ vi.mock("../organisms/create_key_button", () => ({
fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]),
}));
vi.mock("@/app/(dashboard)/hooks/budgets/useBudgetOptions", () => ({
useBudgetOptions: () => ({
data: [
{ budget_id: "svc-a-budget", max_budget: 0.5, created_at: "", updated_at: "" },
{ budget_id: "svc-b-budget", max_budget: 100, created_at: "", updated_at: "" },
],
}),
}));
const routerSettingsMocks = vi.hoisted(() => ({
receivedValue: undefined as { router_settings: Record<string, unknown> } | undefined,
editedValue: null as Record<string, unknown> | null,
@ -182,6 +191,9 @@ describe("KeyEditView", () => {
config: {},
user_id: "default_user_id",
team_id: null,
project_id: null,
key_type: null,
last_active: null,
max_parallel_requests: 10,
metadata: {
logging: [],
@ -1806,6 +1818,86 @@ describe("KeyEditView", () => {
});
});
describe("default customer budget", () => {
const serviceAccountKey = (endUserBudgetId?: string): KeyResponse => ({
...MOCK_KEY_DATA,
metadata: {
service_account_id: "svc-a",
...(endUserBudgetId === undefined ? {} : { end_user_budget_id: endUserBudgetId }),
},
});
const renderEditView = (keyData: KeyResponse, userRole: string = "Admin") => {
const onSubmit = vi.fn().mockResolvedValue(undefined);
renderWithProviders(
<KeyEditView
keyData={keyData}
onCancel={() => {}}
onSubmit={onSubmit}
accessToken={"test-token"}
userID={"test-user"}
userRole={userRole}
premiumUser={false}
/>,
);
return onSubmit;
};
const budgetField = () => screen.findByRole("combobox", { name: "Default Customer Budget" });
const save = async () => userEvent.click(await screen.findByRole("button", { name: /save changes/i }));
it("shows the stored budget and leaves it off an edit that did not touch it", async () => {
const onSubmit = renderEditView(serviceAccountKey("svc-a-budget"));
expect(await budgetField()).toHaveValue("svc-a-budget");
await save();
await waitFor(() => {
expect(onSubmit).toHaveBeenCalled();
});
expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("end_user_budget_id");
});
it("sends the newly chosen budget id", async () => {
const onSubmit = renderEditView(serviceAccountKey());
const user = userEvent.setup();
await chooseSelectOption(user, await budgetField(), /svc-b-budget/);
await save();
await waitFor(() => {
expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ end_user_budget_id: "svc-b-budget" }));
});
});
it("sends an empty string when the stored budget is cleared so the backend removes it", async () => {
const onSubmit = renderEditView(serviceAccountKey("svc-a-budget"));
await budgetField();
await userEvent.click(screen.getByRole("button", { name: "Clear" }));
await save();
await waitFor(() => {
expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ end_user_budget_id: "" }));
});
});
it("keeps the stored budget visible but read-only for a non-admin", async () => {
renderEditView(serviceAccountKey("svc-a-budget"), "Internal User");
const field = await budgetField();
expect(field).toHaveValue("svc-a-budget");
expect(field).toBeDisabled();
});
it("does not render the control on a plain key that has no budget to show", async () => {
renderEditView(MOCK_KEY_DATA);
await screen.findByRole("button", { name: /save changes/i });
expect(screen.queryByRole("combobox", { name: "Default Customer Budget" })).not.toBeInTheDocument();
});
});
describe("estimated output tokens", () => {
const renderEditView = (
keyData: KeyResponse,

View file

@ -47,6 +47,12 @@ import {
toSubmittedValues,
} from "./keyEditFormValues";
import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor";
import { END_USER_BUDGET_HINT, EndUserBudgetSelect } from "../key_team_helpers/EndUserBudgetSelect";
import {
endUserBudgetIdUpdate,
keyOffersEndUserBudget,
storedEndUserBudgetId,
} from "../key_team_helpers/endUserBudgetPayload";
import { ModelMaxBudgetField } from "../key_team_helpers/ModelMaxBudgetEditor";
import { useModelMaxBudgetField } from "../key_team_helpers/useModelMaxBudgetField";
import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
@ -124,8 +130,11 @@ export function KeyEditView({
keyData.budget_fallbacks && typeof keyData.budget_fallbacks === "object" ? keyData.budget_fallbacks : {},
);
const modelBudget = useModelMaxBudgetField(keyData.token, keyData.model_max_budget);
const storedEndUserBudgetIdValue = storedEndUserBudgetId(keyData.metadata);
const [endUserBudgetId, setEndUserBudgetId] = useState<string | null>(storedEndUserBudgetIdValue || null);
const routerSettingsRef = useRef<RouterSettingsAccordionRef>(null);
const keyTypeFieldId = React.useId();
const endUserBudgetFieldId = React.useId();
const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations();
const { data: uiSettingsData } = useUISettings();
const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui);
@ -290,6 +299,11 @@ export function KeyEditView({
modelBudget.applyTo(values);
const endUserBudgetUpdate = endUserBudgetIdUpdate(endUserBudgetId, storedEndUserBudgetIdValue);
if (endUserBudgetUpdate !== undefined) {
values.end_user_budget_id = endUserBudgetUpdate;
}
const routerSettings = routerSettingsUpdate(
routerSettingsRef.current?.getValue()?.router_settings,
keyData.router_settings,
@ -484,6 +498,21 @@ export function KeyEditView({
/>
</Field>
{keyOffersEndUserBudget(keyData.metadata) && (
<Field>
<FieldLabel htmlFor={endUserBudgetFieldId}>
{labelWithHint("Default Customer Budget", END_USER_BUDGET_HINT)}
</FieldLabel>
<EndUserBudgetSelect
id={endUserBudgetFieldId}
accessToken={accessToken}
value={endUserBudgetId}
onChange={setEndUserBudgetId}
canEdit={userRole != null && isProxyAdminRole(userRole)}
/>
</Field>
)}
<KeyRateLimitFields control={form.control} />
<FormField

View file

@ -7711,6 +7711,7 @@ export interface paths {
* - organization_id: Optional[str] - The organization id of the key. If not set, and team_id is set, the organization id will be the same as the team id. If conflict, an error will be raised.
* - project_id: Optional[str] - The project id of the key. When set, models and max_budget are validated against the project's limits.
* - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
* - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Takes precedence over `litellm_settings.max_end_user_budget_id`.
* - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
* - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
* - config: Optional[dict] - any key-specific configs, overrides config in config.yaml
@ -8035,6 +8036,7 @@ export interface paths {
* - team_id: Optional[str] - The team id of the key
* - user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key
* - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
* - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it.
* - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
* - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
* - config: Optional[dict] - any key-specific configs, overrides config in config.yaml
@ -8172,6 +8174,7 @@ export interface paths {
* - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected.
* - organization_id: Optional[str] - The organization id of the key.
* - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
* - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it.
* - models: Optional[list] - Model_name's a user is allowed to call
* - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
* - prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
@ -28507,6 +28510,8 @@ export interface components {
duration?: string | null;
/** Enable Prompt Caching */
enable_prompt_caching?: boolean | null;
/** End User Budget Id */
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Guardrails */
@ -28671,6 +28676,8 @@ export interface components {
duration?: string | null;
/** Enable Prompt Caching */
enable_prompt_caching?: boolean | null;
/** End User Budget Id */
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Expires */
@ -33911,6 +33918,8 @@ export interface components {
duration?: string | null;
/** Enable Prompt Caching */
enable_prompt_caching?: boolean | null;
/** End User Budget Id */
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Expires */
@ -35792,6 +35801,8 @@ export interface components {
duration?: string | null;
/** Enable Prompt Caching */
enable_prompt_caching?: boolean | null;
/** End User Budget Id */
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Grace Period */
@ -39130,6 +39141,8 @@ export interface components {
duration?: string | null;
/** Enable Prompt Caching */
enable_prompt_caching?: boolean | null;
/** End User Budget Id */
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Guardrails */