scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

Merge pull request #40921 from BerriAI/litellm_unified_key_policy_hook

Browse source 
feat(proxy): unified custom_key_policy hook for key generate, update and regenerate
This commit is contained in:
Mateo Wang 2026-09-15 06:34:28 -07:00 • committed by GitHub
commit c8114ba41f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1335 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

276
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -149,6 +149,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequest,
BulkUpdateKeyResponse,
BulkUpdateTeamKeysRequest,
CustomKeyPolicyRequest,
FailedKeyUpdate,
KeySearchWhere,
SuccessfulKeyUpdate,
@ -285,6 +286,7 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
class _CustomKeyHooksModule(Protocol):
user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None
user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None
user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None
def _custom_key_generate_hook(
@ -299,6 +301,161 @@ def _custom_key_update_hook(
return hooks.user_custom_key_update
def _custom_key_policy_hook(
hooks: _CustomKeyHooksModule,
) -> Callable[..., Awaitable[Mapping[str, object]]] | None:
return hooks.user_custom_key_policy
async def _enforce_custom_key_update_policy(
hook: Callable[..., Awaitable[Mapping[str, object]]] | None,
data: UpdateKeyRequest,
) -> None:
if hook is None:
return
if not inspect.iscoroutinefunction(hook):
raise ValueError("user_custom_key_update must be a coroutine")
result: Final = await hook(data)
if not result.get("decision", True):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=result.get("message", "Authentication Failed - Custom Auth Rule"),
)
async def _enforce_custom_key_policy(
hook: Callable[..., Awaitable[Mapping[str, object]]] | None,
build_policy_request: Callable[[], CustomKeyPolicyRequest],
) -> None:
if hook is None:
return
if not inspect.iscoroutinefunction(hook):
raise ValueError("user_custom_key_policy must be a coroutine")
result: Final = await hook(build_policy_request())
if not result.get("decision", True):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=result.get("message", "Authentication Failed - Custom Auth Rule"),
)
_KEY_UPDATE_JSON_STRING_COLUMNS: Final = frozenset({"router_settings", "budget_limits"})
_KEY_METADATA_REQUEST_FIELDS: Final = frozenset(
(*LiteLLM_ManagementEndpoint_MetadataFields_Premium, *LiteLLM_ManagementEndpoint_MetadataFields)
)
def _decode_json_string_column(column: str, value: object) -> object:
if column in _KEY_UPDATE_JSON_STRING_COLUMNS and isinstance(value, str):
return json.loads(value)
return value
def _verification_token_from_row(row: Mapping[str, object]) -> LiteLLM_VerificationToken:
org_id: Final = row["organization_id"] if "organization_id" in row else row.get("org_id")
return LiteLLM_VerificationToken.model_validate(MappingProxyType({**row, "org_id": org_id}))
def _effective_key_after_update(
existing_key_row: LiteLLM_VerificationToken,
non_default_values: Mapping[str, object],
) -> LiteLLM_VerificationToken:
overlay: Final = MappingProxyType(
{column: _decode_json_string_column(column, value) for column, value in non_default_values.items()}
)
return _verification_token_from_row(
MappingProxyType({**existing_key_row.model_dump(), **overlay, "object_permission": None})
)
def _update_policy_request(
operation: Literal["update", "regenerate"],
existing_key_row: LiteLLM_VerificationToken,
non_default_values: Mapping[str, object],
request: UpdateKeyRequest | RegenerateKeyRequest,
) -> CustomKeyPolicyRequest:
return CustomKeyPolicyRequest(
operation=operation,
existing_key=_verification_token_from_row(existing_key_row.model_dump()),
effective_key=_effective_key_after_update(
existing_key_row=existing_key_row, non_default_values=non_default_values
),
request=request,
)
def _generate_budget_windows(
budget_limits: Sequence[BudgetLimitEntry] | None,
) -> tuple[Mapping[str, object], ...] | None:
if not budget_limits:
return None
return tuple(
MappingProxyType(
{
**window.model_dump(),
"reset_at": get_budget_reset_time(budget_duration=window.budget_duration).isoformat(),
}
)
for window in budget_limits
)
def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> LiteLLM_VerificationToken:
requested: Final = data.model_dump(exclude_unset=True, exclude_none=True)
metadata_fields: Final = MappingProxyType(
{field: value for field, value in requested.items() if field in _KEY_METADATA_REQUEST_FIELDS}
)
column_fields: Final = MappingProxyType(
{field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS}
)
metadata: Final = data.metadata or MappingProxyType({})
folded_metadata: Final = {**metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict
columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates in place
expires: Final = (
now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None
)
budget_reset_at: Final = (
get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None
)
key_rotation_at: Final = (
now + timedelta(seconds=duration_in_seconds(duration=data.rotation_interval))
if data.auto_rotate and data.rotation_interval
else None
)
return _verification_token_from_row(
MappingProxyType(
{
**columns,
"metadata": encrypt_callback_vars(folded_metadata),
"expires": expires,
"budget_reset_at": budget_reset_at,
"key_rotation_at": key_rotation_at,
"budget_limits": _generate_budget_windows(data.budget_limits),
"object_permission": None,
}
)
)
_EMPTY_DURATION_MEANS_UNCHANGED: Final = frozenset({"duration", "budget_duration"})
def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None:
changed_fields: Final = MappingProxyType(
{
field: value
for field, value in data.model_dump(exclude_unset=True).items()
if field in UpdateKeyRequest.model_fields
and field != "key"
and not (field in _EMPTY_DURATION_MEANS_UNCHANGED and value == "")
}
)
if not changed_fields:
return None
return UpdateKeyRequest(key=key, **changed_fields)
class _LegacyDumpable(Protocol):
def dict(self) -> Mapping[str, object]: ...
@ -992,6 +1149,7 @@ async def _common_key_generation_helper(
litellm_changed_by: str | None,
team_table: LiteLLM_TeamTableCachedObj | None,
) -> GenerateKeyResponse:
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
llm_router,
@ -1140,6 +1298,16 @@ async def _common_key_generation_helper(
"litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e
)
await _enforce_custom_key_policy(
hook=_custom_key_policy_hook(proxy_server),
build_policy_request=lambda: CustomKeyPolicyRequest(
operation="generate",
existing_key=None,
effective_key=_effective_key_for_generate(data=data, now=datetime.now(timezone.utc)),
request=data,
),
)
# TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable
_budget_id = data.budget_id
if prisma_client is not None and data.soft_budget is not None:
@ -2325,12 +2493,6 @@ async def prepare_key_update_data(
# sentinel for Json? columns, so store the JSON literal null
non_default_values["budget_limits"] = json.dumps(None)
if "object_permission" in non_default_values:
non_default_values = await _handle_update_object_permission(
data_json=non_default_values,
existing_key_row=existing_key_row,
)
_metadata: Final = existing_key_row.metadata or {}
# validate model_max_budget
@ -2351,13 +2513,12 @@ async def prepare_key_update_data(
async def _handle_update_object_permission(
data_json: dict,
existing_key_row: LiteLLM_VerificationToken,
prisma_client: PrismaClient,
) -> dict:
"""
Handle the update of object permission.
"""
from litellm.proxy.proxy_server import prisma_client
"""Persist the requested object permission row and swap it for its id, only after the key policy allowed the write."""
if "object_permission" not in data_json:
return data_json
# Use the common helper to handle the object permission update
object_permission_id: Final = await handle_update_object_permission_common(
data_json=data_json,
existing_object_permission_id=existing_key_row.object_permission_id,
@ -2491,6 +2652,7 @@ async def _process_single_key_update(
llm_router: Router | None,
user_custom_key_update: Callable | None = None,
existing_key_row: LiteLLM_VerificationToken | None = None,
user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None = None,
) -> dict[str, object]:
"""
Process a single key update with all validations and checks.
@ -2603,6 +2765,16 @@ async def _process_single_key_update(
data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
)
await _enforce_custom_key_policy(
hook=user_custom_key_policy,
build_policy_request=lambda: _update_policy_request(
operation="update",
existing_key_row=existing_key_row,
non_default_values=non_default_values,
request=update_key_request,
),
)
# Update key in database
if prisma_client is None:
raise HTTPException(
@ -2610,7 +2782,12 @@ async def _process_single_key_update(
detail={"error": "Database not connected"},
)
_data: Final = {**non_default_values, "token": update_key_request.key}
update_values: Final = await _handle_update_object_permission(
data_json=non_default_values,
existing_key_row=existing_key_row,
prisma_client=prisma_client,
)
_data: Final = {**update_values, "token": update_key_request.key}
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
"Mapping[str, object] | None",
await prisma_client.update_data(token=update_key_request.key, data=_data),
@ -3103,19 +3280,7 @@ async def update_key_fn(
user_api_key_cache=user_api_key_cache,
)
# Custom key update hook
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook(
proxy_server
)
if custom_key_update_hook is not None:
if inspect.iscoroutinefunction(custom_key_update_hook):
result: Final = await custom_key_update_hook(data)
else:
raise ValueError("user_custom_key_update must be a coroutine")
decision: Final = result.get("decision", True)
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
if not decision:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=data)
# Enforce upperbound key params on update (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
@ -3142,21 +3307,36 @@ async def update_key_fn(
existing_key_alias=existing_key_row.key_alias,
)
await _enforce_custom_key_policy(
hook=_custom_key_policy_hook(proxy_server),
build_policy_request=lambda: _update_policy_request(
operation="update",
existing_key_row=existing_key_row,
non_default_values=non_default_values,
request=data,
),
)
if prisma_client is None:
raise Exception("Not connected to DB!")
update_values: Final = await _handle_update_object_permission(
data_json=non_default_values,
existing_key_row=existing_key_row,
prisma_client=prisma_client,
)
changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
response: Final = (
await _update_key_row_with_soft_budget(
prisma_client=prisma_client,
key=key,
data=data,
non_default_values=non_default_values,
non_default_values=update_values,
existing_key_row=existing_key_row,
changed_by=changed_by,
)
if "soft_budget" in data.model_fields_set
else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key}))
else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key}))
)
# Delete - key from cache, since it's been updated!
@ -3291,6 +3471,7 @@ async def bulk_update_keys(
)
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
@ -3338,6 +3519,7 @@ async def bulk_update_keys(
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=custom_key_update_hook,
user_custom_key_policy=custom_key_policy_hook,
)
successful_updates.append(
@ -3455,6 +3637,7 @@ async def bulk_update_team_keys(
)
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server)
if prisma_client is None:
raise HTTPException(
@ -3585,6 +3768,7 @@ async def bulk_update_team_keys(
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=custom_key_update_hook,
user_custom_key_policy=custom_key_policy_hook,
existing_key_row=existing_by_token[db_token],
)
@ -5118,6 +5302,7 @@ async def _execute_virtual_key_regeneration(
proxy_logging_obj: ProxyLogging,
) -> GenerateKeyResponse:
"""Generate new token, update DB, invalidate cache, and return response."""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import hash_token
# Mirror the /key/update ownership rebind guard. See helper docstring.
@ -5165,6 +5350,9 @@ async def _execute_virtual_key_regeneration(
non_default_values = {}
if data is not None:
update_request: Final = _regenerate_request_as_update_request(key=hashed_api_key, data=data)
if update_request is not None:
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request)
# Enforce upperbound key params on regenerate (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
non_default_values = await prepare_key_update_data(
@ -5175,7 +5363,21 @@ async def _execute_virtual_key_regeneration(
if new_key_alias != key_in_db.key_alias:
_validate_key_alias_format(key_alias=new_key_alias)
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
update_data.update(non_default_values)
await _enforce_custom_key_policy(
hook=_custom_key_policy_hook(proxy_server),
build_policy_request=lambda: _update_policy_request(
operation="regenerate",
existing_key_row=key_in_db,
non_default_values=non_default_values,
request=data if data is not None else RegenerateKeyRequest(),
),
)
update_values: Final = await _handle_update_object_permission(
data_json=non_default_values,
existing_key_row=key_in_db,
prisma_client=prisma_client,
)
update_data.update(update_values)
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
# Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash,
@ -5185,6 +5387,13 @@ async def _execute_virtual_key_regeneration(
prisma_client=prisma_client,
)
await _persist_deleted_verification_tokens(
keys=[key_in_db],
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
# If grace period set, insert deprecated key so old key remains valid
await _insert_deprecated_key(
prisma_client=prisma_client,
@ -5484,17 +5693,6 @@ async def regenerate_key_fn(
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
litellm_changed_by = None
# Save the old key record to deleted table before regeneration.
# This preserves key_alias and team_id metadata for historical spend records.
# If this fails, abort the regeneration to avoid permanently losing the
# old hash→metadata mapping.
await _persist_deleted_verification_tokens(
keys=[_key_in_db],
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
return await _execute_virtual_key_regeneration(
prisma_client=prisma_client,
llm_router=llm_router,

9
litellm/proxy/proxy_server.py
View file

@ -928,6 +928,7 @@ def cleanup_router_config_variables():
user_custom_auth_path, \
user_custom_key_generate, \
user_custom_key_update, \
user_custom_key_policy, \
user_custom_sso, \
user_custom_ui_sso_sign_in_handler, \
use_background_health_checks, \
@ -945,6 +946,7 @@ def cleanup_router_config_variables():
user_custom_auth_path = None
user_custom_key_generate = None
user_custom_key_update = None
user_custom_key_policy = None
TEAM_METADATA_VALIDATOR_REGISTRY.set(None)
TEAM_METADATA_SCHEMA_REGISTRY.set(())
user_custom_sso = None
@ -2369,6 +2371,7 @@ user_custom_key_generate = None
_pkce_no_redis_warning_emitted: bool = False
_cp_no_redis_warning_emitted: bool = False
user_custom_key_update = None
user_custom_key_policy = None
user_custom_sso = None
user_custom_ui_sso_sign_in_handler = None
use_background_health_checks = None
@ -4256,6 +4259,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Final[dict[str, tuple[str, ...]]] = {
"custom_auth",
"custom_key_generate",
"custom_key_update",
"custom_key_policy",
"custom_team_metadata_validate",
"custom_sso",
"custom_ui_sso_sign_in_handler",
@ -5405,6 +5409,7 @@ class ProxyConfig:
user_custom_auth_path, \
user_custom_key_generate, \
user_custom_key_update, \
user_custom_key_policy, \
user_custom_sso, \
user_custom_ui_sso_sign_in_handler, \
use_background_health_checks, \
@ -5942,6 +5947,10 @@ class ProxyConfig:
if custom_key_update is not None:
user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path)
custom_key_policy: Final = general_settings.get("custom_key_policy", None)
if custom_key_policy is not None:
user_custom_key_policy = get_instance_fn(value=custom_key_policy, config_file_path=config_file_path)
custom_team_metadata_validate: Final = general_settings.get("custom_team_metadata_validate", None)
TEAM_METADATA_VALIDATOR_REGISTRY.set(
get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path)

26
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -1,9 +1,12 @@
from datetime import datetime
from typing import Any, Final, Literal
from typing import Any, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, model_validator
from typing_extensions import ReadOnly, TypedDict
from litellm.models.verification_token import LiteLLM_VerificationToken
from litellm.proxy._types import GenerateKeyRequest, RegenerateKeyRequest, UpdateKeyRequest
from litellm.types.llms.base import LiteLLMPydanticObjectBase
from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
@ -123,3 +126,24 @@ class BulkUpdateTeamKeysRequest(BaseModel):
if not has_key_ids and not self.all_keys_in_team:
raise ValueError("Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`.")
return self
CustomKeyPolicyOperation: TypeAlias = Literal["generate", "update", "regenerate"]
class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase):
"""What `general_settings.custom_key_policy` receives.
`effective_key` is the verification token row as it will be written: the existing row overlaid with the
requested changes, with `duration` resolved to `expires` and `budget_duration` to `budget_reset_at`. Values the
proxy fills in after the policy stay at their defaults: `token`, `key_name`, `created_by`, `updated_by` and the
soft-budget `budget_id` on generate, the rotated token on regenerate, and the `object_permission` relation on
every operation (`object_permission_id` is set; read `request.object_permission` for the requested change).
"""
model_config = ConfigDict(protected_namespaces=(), frozen=True)
operation: CustomKeyPolicyOperation
existing_key: LiteLLM_VerificationToken | None
effective_key: LiteLLM_VerificationToken
request: GenerateKeyRequest | UpdateKeyRequest | RegenerateKeyRequest

1077
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

File diff suppressed because it is too large Load diff

1
tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py
View file

@ -43,6 +43,7 @@ def test_litellm_settings_callback_list_strips_remote_urls(field):
"custom_auth",
"custom_key_generate",
"custom_key_update",
"custom_key_policy",
"custom_sso",
"custom_ui_sso_sign_in_handler",
],