mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
feat(proxy): per-key default budget for dynamically created customers
A service-account key can now carry end_user_budget_id in its metadata. When a request through that key names a customer that does not exist yet, the key's budget is applied to the new customer from the first request and wins over the proxy-wide max_end_user_budget_id. A customer with an explicitly assigned budget keeps it. The Admin UI exposes the setting on service-account key creation and key edit, and only proxy admins may set or clear it. Resolves LIT-7996 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
0f5bc0ffa9
commit
cc11653152
18 changed files with 1235 additions and 62 deletions
|
|
@ -1215,6 +1215,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
|
||||
|
|
@ -4721,6 +4722,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 = [
|
||||
|
|
|
|||
|
|
@ -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,13 +1403,11 @@ 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
|
||||
)
|
||||
verbose_proxy_logger.warning("Default end user budget not found in database: %s", default_budget_id)
|
||||
return None
|
||||
|
||||
_budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict())
|
||||
|
|
@ -1469,47 +1482,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 +1761,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 +1769,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 +1798,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 +1816,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 +1835,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 +1862,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 +1902,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 +1916,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 +1932,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 +1955,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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -706,6 +708,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
|
|||
valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"]
|
||||
if end_user_params.get("end_user_tpd_limit") is not None:
|
||||
valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"]
|
||||
if end_user_params.get("end_user_max_budget") is not None:
|
||||
valid_token.end_user_max_budget = end_user_params["end_user_max_budget"]
|
||||
if end_user_params.get("allowed_model_region") is not None:
|
||||
valid_token.allowed_model_region = end_user_params["allowed_model_region"]
|
||||
if end_user_params.get("end_user_model_max_budget") is not None:
|
||||
|
|
@ -2680,6 +2684,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 +2695,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 +2761,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 +2866,16 @@ 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,
|
||||
)
|
||||
|
||||
skip_budget_checks: Final = _should_skip_budget_checks(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
|
|
@ -2945,6 +2964,37 @@ 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,
|
||||
) -> 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.
|
||||
The budget replaces the proxy-wide one wholesale: a key budget with no cap also lifts the cap."""
|
||||
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
|
||||
|
||||
valid_token.end_user_max_budget = default_budget.max_budget
|
||||
valid_token.end_user_tpm_limit = default_budget.tpm_limit
|
||||
valid_token.end_user_rpm_limit = default_budget.rpm_limit
|
||||
valid_token.end_user_tpd_limit = default_budget.tpd_limit
|
||||
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 +3144,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 +3422,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 +3432,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 +3448,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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -2887,6 +2895,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 +3037,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,
|
||||
|
|
@ -5383,6 +5434,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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -222,6 +222,59 @@ 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
|
||||
|
||||
|
||||
def test_end_user_budget_max_budget_reaches_the_token():
|
||||
from litellm.proxy.auth.user_api_key_auth import _apply_budget_limits_to_end_user_params
|
||||
|
||||
end_user_params = {"end_user_id": "user_1"}
|
||||
_apply_budget_limits_to_end_user_params(
|
||||
end_user_params=end_user_params,
|
||||
budget_info=LiteLLM_BudgetTable(max_budget=20.0),
|
||||
end_user_id="user_1",
|
||||
)
|
||||
result = update_valid_token_with_end_user_params(UserAPIKeyAuth(token="test_token"), end_user_params)
|
||||
|
||||
assert result.end_user_max_budget == 20.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,
|
||||
|
|
|
|||
|
|
@ -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,142 @@ 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,
|
||||
) -> UserAPIKeyAuth:
|
||||
"""Run the centralized checks with a fake DB and return the token handed to budget reservation."""
|
||||
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=None),
|
||||
"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
|
||||
|
||||
|
||||
class _RecordingTeamModelBudgetLimiter:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
};
|
||||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
@ -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}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
|
@ -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("");
|
||||
});
|
||||
});
|
||||
|
|
@ -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;
|
||||
};
|
||||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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={
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
10
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
10
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -28334,6 +28334,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 */
|
||||
|
|
@ -28498,6 +28500,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 */
|
||||
|
|
@ -33700,6 +33704,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 */
|
||||
|
|
@ -35571,6 +35577,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 */
|
||||
|
|
@ -38851,6 +38859,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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue