mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
feat(proxy): unified custom_key_policy hook for key generate, update and regenerate
Adds general_settings.custom_key_policy, one coroutine that receives the
operation ("generate", "update", "regenerate"), the existing key row, the
effective row as it will be written, and the raw request, and can deny with a
403. It runs after the request has been normalized and before the first DB
write on /key/generate, /key/service-account/generate, /key/update,
/key/bulk_update, /team/key/bulk_update and /key/{key}/regenerate. The two
legacy hooks keep running unchanged on the raw request.
This commit is contained in:
mateo-berri
parent
c9a3c5a414
commit
577a4e94aa
5 changed files with 808 additions and 0 deletions
|
|
@ -148,6 +148,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
|||
BulkUpdateKeyRequest,
|
||||
BulkUpdateKeyResponse,
|
||||
BulkUpdateTeamKeysRequest,
|
||||
CustomKeyPolicyRequest,
|
||||
FailedKeyUpdate,
|
||||
KeySearchWhere,
|
||||
SuccessfulKeyUpdate,
|
||||
|
|
@ -280,6 +281,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(
|
||||
|
|
@ -294,6 +296,12 @@ 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,
|
||||
|
|
@ -310,6 +318,113 @@ async def _enforce_custom_key_update_policy(
|
|||
)
|
||||
|
||||
|
||||
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}))
|
||||
|
||||
|
||||
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 budget_limits is None:
|
||||
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}
|
||||
)
|
||||
request_metadata: Final = data.metadata or MappingProxyType({})
|
||||
folded_metadata: Final = {**request_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 data_json 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
|
||||
)
|
||||
return _verification_token_from_row(
|
||||
MappingProxyType(
|
||||
{
|
||||
**columns,
|
||||
"metadata": encrypt_callback_vars(folded_metadata),
|
||||
"expires": expires,
|
||||
"budget_reset_at": budget_reset_at,
|
||||
"budget_limits": _generate_budget_windows(data.budget_limits),
|
||||
"object_permission": None,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None:
|
||||
changed_fields: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -1016,6 +1131,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,
|
||||
|
|
@ -1164,6 +1280,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:
|
||||
|
|
@ -2496,6 +2622,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.
|
||||
|
|
@ -2606,6 +2733,16 @@ async def _process_single_key_update(
|
|||
# Prepare update data
|
||||
non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row)
|
||||
|
||||
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(
|
||||
|
|
@ -3112,6 +3249,16 @@ async def update_key_fn(
|
|||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row)
|
||||
|
||||
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,
|
||||
),
|
||||
)
|
||||
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias: Final = non_default_values.get("key_alias", None)
|
||||
if new_key_alias != existing_key_row.key_alias:
|
||||
|
|
@ -3280,6 +3427,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(
|
||||
|
|
@ -3327,6 +3475,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(
|
||||
|
|
@ -3444,6 +3593,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(
|
||||
|
|
@ -3574,6 +3724,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],
|
||||
)
|
||||
|
||||
|
|
@ -5156,6 +5307,15 @@ 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)
|
||||
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_data.update(non_default_values)
|
||||
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
|
||||
|
||||
|
|
|
|||
|
|
@ -924,6 +924,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, \
|
||||
|
|
@ -941,6 +942,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
|
||||
|
|
@ -2365,6 +2367,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
|
||||
|
|
@ -4250,6 +4253,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",
|
||||
|
|
@ -5403,6 +5407,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, \
|
||||
|
|
@ -5940,6 +5945,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)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,9 @@ from typing import Any, Final, Literal
|
|||
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,17 @@ 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 = Literal["generate", "update", "regenerate"]
|
||||
|
||||
|
||||
class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase):
|
||||
"""What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request."""
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=(), frozen=True)
|
||||
|
||||
operation: CustomKeyPolicyOperation
|
||||
existing_key: LiteLLM_VerificationToken | None
|
||||
effective_key: LiteLLM_VerificationToken
|
||||
request: GenerateKeyRequest | UpdateKeyRequest | RegenerateKeyRequest
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from contextlib import ExitStack
|
||||
from typing import Final
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.proxy._types import (
|
|||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_VerificationToken,
|
||||
LiteLLMKeyType,
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
ProxyException,
|
||||
|
|
@ -38,6 +40,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
_check_project_key_limits,
|
||||
_check_team_key_limits,
|
||||
_common_key_generation_helper,
|
||||
_effective_key_after_update,
|
||||
_effective_key_for_generate,
|
||||
_enforce_custom_key_policy,
|
||||
_enforce_upperbound_key_params,
|
||||
_execute_virtual_key_regeneration,
|
||||
_get_and_validate_existing_key,
|
||||
|
|
@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
validate_key_team_change,
|
||||
)
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import CustomKeyPolicyRequest
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
|
|
@ -12178,6 +12184,586 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit
|
|||
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
|
||||
|
||||
|
||||
_POLICY_DENIAL_MESSAGE = "key duration must be 7d or less"
|
||||
_POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b"
|
||||
_POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None}
|
||||
|
||||
|
||||
def _seven_day_policy(received: list[CustomKeyPolicyRequest]):
|
||||
async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
|
||||
received.append(policy_request)
|
||||
expires = policy_request.effective_key.expires
|
||||
if isinstance(expires, datetime) and expires > datetime.now(timezone.utc) + timedelta(days=7):
|
||||
return {"decision": False, "message": _POLICY_DENIAL_MESSAGE}
|
||||
return {"decision": True}
|
||||
|
||||
return policy
|
||||
|
||||
|
||||
def _assert_expires_in(effective_key: LiteLLM_VerificationToken, duration: str) -> None:
|
||||
expires = effective_key.expires
|
||||
assert isinstance(expires, datetime)
|
||||
assert expires.tzinfo is not None
|
||||
expected = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration=duration))
|
||||
assert abs((expires - expected).total_seconds()) < 60
|
||||
|
||||
|
||||
def _regenerate_policy_mocks(policy, insert_deprecated_key: AsyncMock, persist: AsyncMock) -> ExitStack:
|
||||
stack = ExitStack()
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: deterministic token setup for the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value="sk-newtoken1234ab12",
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: grace-period write must not run on a denied regenerate
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
|
||||
insert_deprecated_key,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: archival write must not run on a denied regenerate
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens",
|
||||
persist,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: cache eviction is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: rotation callback is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook
|
||||
)
|
||||
return stack
|
||||
|
||||
|
||||
async def _regenerate_under_policy(mock_prisma_client, existing_key, data):
|
||||
return await _execute_virtual_key_regeneration(
|
||||
prisma_client=mock_prisma_client,
|
||||
key_in_db=existing_key,
|
||||
hashed_api_key="abc123",
|
||||
key="abc123",
|
||||
data=data,
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_rejects_when_custom_key_policy_denies_the_effective_expiry():
|
||||
existing_key = _make_regenerate_existing_key()
|
||||
mock_prisma_client = _make_regenerate_mock_prisma()
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
insert_deprecated_key = AsyncMock()
|
||||
persist = AsyncMock()
|
||||
|
||||
with _regenerate_policy_mocks(_seven_day_policy(received), insert_deprecated_key, persist):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="3000d"))
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE
|
||||
insert_deprecated_key.assert_not_awaited()
|
||||
persist.assert_not_awaited()
|
||||
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0
|
||||
assert [policy_request.operation for policy_request in received] == ["regenerate"]
|
||||
assert received[0].existing_key is not None
|
||||
assert received[0].existing_key.token == "abc123"
|
||||
assert isinstance(received[0].request, RegenerateKeyRequest)
|
||||
assert received[0].request.duration == "3000d"
|
||||
_assert_expires_in(received[0].effective_key, "3000d")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_within_custom_key_policy_rotates_the_key():
|
||||
existing_key = _make_regenerate_existing_key()
|
||||
mock_prisma_client = _make_regenerate_mock_prisma()
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
persist = AsyncMock()
|
||||
|
||||
with _regenerate_policy_mocks(_seven_day_policy(received), AsyncMock(), persist):
|
||||
await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="5d"))
|
||||
|
||||
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
|
||||
persist.assert_awaited_once()
|
||||
assert persist.call_args.kwargs["keys"] == [existing_key]
|
||||
assert [policy_request.operation for policy_request in received] == ["regenerate"]
|
||||
_assert_expires_in(received[0].effective_key, "5d")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()])
|
||||
async def test_regenerate_without_changes_still_runs_custom_key_policy(data):
|
||||
existing_key = _make_regenerate_existing_key()
|
||||
mock_prisma_client = _make_regenerate_mock_prisma()
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
|
||||
async def freeze_rotation(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
|
||||
received.append(policy_request)
|
||||
return {"decision": False, "message": "key rotation is frozen"}
|
||||
|
||||
with _regenerate_policy_mocks(freeze_rotation, AsyncMock(), AsyncMock()):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _regenerate_under_policy(mock_prisma_client, existing_key, data)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "key rotation is frozen"
|
||||
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0
|
||||
assert [policy_request.operation for policy_request in received] == ["regenerate"]
|
||||
assert received[0].existing_key == existing_key
|
||||
assert received[0].effective_key == existing_key
|
||||
|
||||
|
||||
def _policy_existing_team_key() -> LiteLLM_VerificationToken:
|
||||
return LiteLLM_VerificationToken(
|
||||
token=_POLICY_HASHED_TOKEN, user_id="test-user", team_id="team-a", max_budget=200.0
|
||||
)
|
||||
|
||||
|
||||
def _setup_update_key_fn_policy_mocks(monkeypatch, existing_key: LiteLLM_VerificationToken) -> AsyncMock:
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None)
|
||||
mock_prisma_client.update_data = AsyncMock(return_value={"data": {"max_budget": 50.0, "team_id": "team-a"}})
|
||||
_setup_update_key_mocks(monkeypatch, mock_prisma_client)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", AsyncMock(return_value=None)
|
||||
)
|
||||
return mock_prisma_client
|
||||
|
||||
|
||||
def _assert_update_policy_request(policy_request: CustomKeyPolicyRequest, request: UpdateKeyRequest) -> None:
|
||||
assert policy_request.operation == "update"
|
||||
assert policy_request.request is request
|
||||
assert policy_request.existing_key is not None
|
||||
assert policy_request.existing_key.max_budget == 200.0
|
||||
assert policy_request.effective_key.team_id == "team-a"
|
||||
assert policy_request.effective_key.user_id == "test-user"
|
||||
assert policy_request.effective_key.max_budget == 50.0
|
||||
_assert_expires_in(policy_request.effective_key, request.duration or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeypatch):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn
|
||||
|
||||
mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key())
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
policy = _seven_day_policy(received)
|
||||
data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: cache eviction is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_policy", policy), # test-quality-ok: inject policy hook
|
||||
):
|
||||
await update_key_fn(
|
||||
request=MagicMock(),
|
||||
data=data,
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
mock_prisma_client.update_data.assert_awaited_once()
|
||||
assert len(received) == 1
|
||||
_assert_update_policy_request(received[0], data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_fn_rejects_when_custom_key_policy_denies(monkeypatch):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn
|
||||
|
||||
mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key())
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
policy = _seven_day_policy(received)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.user_custom_key_policy", policy): # test-quality-ok: inject policy hook
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await update_key_fn(
|
||||
request=MagicMock(),
|
||||
data=UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0),
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert exc_info.value.message == _POLICY_DENIAL_MESSAGE
|
||||
mock_prisma_client.update_data.assert_not_awaited()
|
||||
assert [policy_request.operation for policy_request in received] == ["update"]
|
||||
_assert_expires_in(received[0].effective_key, "3000d")
|
||||
|
||||
|
||||
async def _process_single_key_update_under_policy(prisma_client: AsyncMock, data: UpdateKeyRequest, policy):
|
||||
with (
|
||||
patch( # test-quality-ok: cache eviction is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch( # test-quality-ok: update callback is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
return await _process_single_key_update(
|
||||
update_key_request=data,
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
llm_router=None,
|
||||
existing_key_row=_policy_existing_team_key(),
|
||||
user_custom_key_policy=policy,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_single_key_update_runs_custom_key_policy_on_the_effective_row():
|
||||
mock_prisma_client = AsyncMock()
|
||||
updated_row = MagicMock()
|
||||
updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"}
|
||||
mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row})
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0)
|
||||
|
||||
result = await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received))
|
||||
|
||||
assert result["max_budget"] == 50.0
|
||||
mock_prisma_client.update_data.assert_awaited_once()
|
||||
assert len(received) == 1
|
||||
_assert_update_policy_request(received[0], data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_single_key_update_rejects_when_custom_key_policy_denies():
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.update_data = AsyncMock()
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received))
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE
|
||||
mock_prisma_client.update_data.assert_not_awaited()
|
||||
assert [policy_request.operation for policy_request in received] == ["update"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateKeyRequest,
|
||||
BulkUpdateKeyRequestItem,
|
||||
)
|
||||
|
||||
existing_keys = [
|
||||
LiteLLM_VerificationToken(token="test-key-1", user_id="user-123", max_budget=None),
|
||||
LiteLLM_VerificationToken(token="test-key-2", user_id="user-123", max_budget=50.0),
|
||||
]
|
||||
updated_row = MagicMock()
|
||||
updated_row.model_dump.return_value = {"user_id": "user-123", "max_budget": 100.0}
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=existing_keys)
|
||||
mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row})
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
|
||||
async def cap_max_budget(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
|
||||
received.append(policy_request)
|
||||
max_budget = policy_request.effective_key.max_budget
|
||||
if max_budget is not None and max_budget > 100:
|
||||
return {"decision": False, "message": "max_budget must be 100 or less"}
|
||||
return {"decision": True}
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", cap_max_budget)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: cache eviction is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch( # test-quality-ok: update callback is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
response = await bulk_update_keys(
|
||||
data=BulkUpdateKeyRequest(
|
||||
keys=[
|
||||
BulkUpdateKeyRequestItem(key="test-key-1", max_budget=100.0),
|
||||
BulkUpdateKeyRequestItem(key="test-key-2", max_budget=500.0),
|
||||
]
|
||||
),
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert [update.key for update in response.successful_updates] == ["test-key-1"]
|
||||
assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [
|
||||
("test-key-2", "max_budget must be 100 or less")
|
||||
]
|
||||
assert mock_prisma_client.update_data.await_count == 1
|
||||
assert [policy_request.operation for policy_request in received] == ["update", "update"]
|
||||
assert [policy_request.effective_key.max_budget for policy_request in received] == [100.0, 500.0]
|
||||
assert [
|
||||
policy_request.existing_key.max_budget if policy_request.existing_key is not None else "missing"
|
||||
for policy_request in received
|
||||
] == [None, 50.0]
|
||||
|
||||
|
||||
def _policy_generate_prisma() -> MagicMock:
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1"))
|
||||
mock_prisma.jsonify_object = MagicMock(side_effect=lambda data: json.loads(data) if isinstance(data, str) else data)
|
||||
return mock_prisma
|
||||
|
||||
|
||||
def _generate_policy_mocks(mock_prisma: MagicMock, generate_key_helper: AsyncMock, policy) -> ExitStack:
|
||||
stack = ExitStack()
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)) # test-quality-ok: fake DB
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", None)) # test-quality-ok: no router in test
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) # test-quality-ok: premium fields
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")) # test-quality-ok: admin
|
||||
stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())) # test-quality-ok: cache
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: the key write must not run on a denied generate
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
generate_key_helper,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook
|
||||
)
|
||||
return stack
|
||||
|
||||
|
||||
def _generate_request(duration: str, organization_id: str | None) -> GenerateKeyRequest:
|
||||
return GenerateKeyRequest(
|
||||
duration=duration,
|
||||
organization_id=organization_id,
|
||||
guardrails=["g1"],
|
||||
tags=["t1"],
|
||||
soft_budget=10.0,
|
||||
max_budget=20.0,
|
||||
)
|
||||
|
||||
|
||||
def _assert_generate_policy_request(
|
||||
policy_request: CustomKeyPolicyRequest, duration: str, organization_id: str | None
|
||||
) -> None:
|
||||
assert policy_request.operation == "generate"
|
||||
assert policy_request.existing_key is None
|
||||
assert policy_request.effective_key.org_id == organization_id
|
||||
assert policy_request.effective_key.max_budget == 20.0
|
||||
assert policy_request.effective_key.metadata["guardrails"] == ["g1"]
|
||||
assert policy_request.effective_key.metadata["tags"] == ["t1"]
|
||||
_assert_expires_in(policy_request.effective_key, duration)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_key_rejects_when_custom_key_policy_denies_before_any_write():
|
||||
mock_prisma = _policy_generate_prisma()
|
||||
generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY)
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
data = _generate_request("3000d", organization_id="org-1")
|
||||
|
||||
with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await generate_key_fn(
|
||||
data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None
|
||||
)
|
||||
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert exc_info.value.message == _POLICY_DENIAL_MESSAGE
|
||||
mock_prisma.db.litellm_budgettable.create.assert_not_awaited()
|
||||
generate_key_helper.assert_not_awaited()
|
||||
assert len(received) == 1
|
||||
_assert_generate_policy_request(received[0], "3000d", organization_id="org-1")
|
||||
assert received[0].request is data
|
||||
assert data.duration == "3000d"
|
||||
assert data.guardrails == ["g1"]
|
||||
assert data.tags == ["t1"]
|
||||
assert data.organization_id == "org-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_key_within_custom_key_policy_creates_the_key():
|
||||
mock_prisma = _policy_generate_prisma()
|
||||
generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY)
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
data = _generate_request("5d", organization_id=None)
|
||||
|
||||
with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)):
|
||||
await generate_key_fn(
|
||||
data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_budgettable.create.assert_awaited_once()
|
||||
generate_key_helper.assert_awaited_once()
|
||||
assert len(received) == 1
|
||||
_assert_generate_policy_request(received[0], "5d", organization_id=None)
|
||||
assert received[0].request is data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_account_generate_rejects_when_custom_key_policy_denies():
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import generate_service_account_key_fn
|
||||
|
||||
mock_prisma = _policy_generate_prisma()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock())
|
||||
generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY)
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
|
||||
with (
|
||||
_generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)),
|
||||
patch( # test-quality-ok: team lookup is outside the policy path
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await generate_service_account_key_fn(
|
||||
data=GenerateKeyRequest(team_id="team-1", duration="3000d"),
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE
|
||||
generate_key_helper.assert_not_awaited()
|
||||
mock_prisma.db.litellm_budgettable.create.assert_not_awaited()
|
||||
assert [policy_request.operation for policy_request in received] == ["generate"]
|
||||
assert received[0].existing_key is None
|
||||
assert received[0].effective_key.team_id == "team-1"
|
||||
assert received[0].effective_key.user_id is None
|
||||
_assert_expires_in(received[0].effective_key, "3000d")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_effective_key_after_update_decodes_json_string_columns_and_keeps_omitted_fields():
|
||||
existing_key = LiteLLM_VerificationToken(token="tok", user_id="u1", team_id="team-a")
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=UpdateKeyRequest(
|
||||
key="tok", router_settings={"num_retries": 3}, budget_limits=[{"budget_duration": "1d", "max_budget": 2.0}]
|
||||
),
|
||||
existing_key_row=existing_key,
|
||||
)
|
||||
assert isinstance(non_default_values["router_settings"], str)
|
||||
assert isinstance(non_default_values["budget_limits"], str)
|
||||
|
||||
effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values)
|
||||
|
||||
assert effective_key.router_settings == {"num_retries": 3}
|
||||
assert effective_key.budget_limits is not None
|
||||
assert effective_key.budget_limits[0]["max_budget"] == 2.0
|
||||
assert effective_key.budget_limits[0]["budget_duration"] == "1d"
|
||||
assert effective_key.budget_limits[0]["reset_at"] is not None
|
||||
assert effective_key.team_id == "team-a"
|
||||
assert effective_key.user_id == "u1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration():
|
||||
existing_key = LiteLLM_VerificationToken(token="tok", expires=datetime(2027, 1, 1, tzinfo=timezone.utc))
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=UpdateKeyRequest(key="tok", duration="-1"), existing_key_row=existing_key
|
||||
)
|
||||
|
||||
effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values)
|
||||
|
||||
assert effective_key.expires is None
|
||||
|
||||
|
||||
def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it():
|
||||
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
data = GenerateKeyRequest(
|
||||
duration="5d",
|
||||
organization_id="org-1",
|
||||
metadata={"a": 1},
|
||||
guardrails=["g1"],
|
||||
tags=["t1"],
|
||||
budget_duration="1d",
|
||||
max_budget=3.0,
|
||||
key_type=LiteLLMKeyType.LLM_API,
|
||||
)
|
||||
|
||||
effective_key = _effective_key_for_generate(data=data, now=now)
|
||||
|
||||
assert effective_key.expires == now + timedelta(days=5)
|
||||
assert effective_key.org_id == "org-1"
|
||||
assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]}
|
||||
assert effective_key.max_budget == 3.0
|
||||
assert effective_key.budget_duration == "1d"
|
||||
assert effective_key.budget_reset_at is not None
|
||||
assert effective_key.key_type == "llm_api"
|
||||
assert effective_key.allowed_routes == ["llm_api_routes"]
|
||||
assert data.metadata == {"a": 1}
|
||||
assert data.guardrails == ["g1"]
|
||||
assert data.tags == ["t1"]
|
||||
assert data.duration == "5d"
|
||||
|
||||
|
||||
def test_effective_key_for_generate_without_duration_never_expires():
|
||||
effective_key = _effective_key_for_generate(
|
||||
data=GenerateKeyRequest(), now=datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
)
|
||||
|
||||
assert effective_key.expires is None
|
||||
assert effective_key.budget_reset_at is None
|
||||
assert effective_key.key_type == "default"
|
||||
|
||||
|
||||
def _policy_request_for_generate() -> CustomKeyPolicyRequest:
|
||||
return CustomKeyPolicyRequest(
|
||||
operation="generate",
|
||||
existing_key=None,
|
||||
effective_key=LiteLLM_VerificationToken(token="tok"),
|
||||
request=GenerateKeyRequest(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_custom_key_policy_rejects_a_sync_hook():
|
||||
def sync_hook(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
|
||||
return {"decision": True}
|
||||
|
||||
with pytest.raises(ValueError, match="user_custom_key_policy must be a coroutine"):
|
||||
await _enforce_custom_key_policy(hook=sync_hook, build_policy_request=_policy_request_for_generate)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_custom_key_policy_uses_the_default_denial_message():
|
||||
async def deny(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
|
||||
return {"decision": False}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _enforce_custom_key_policy(hook=deny, build_policy_request=_policy_request_for_generate)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token():
|
||||
"""
|
||||
|
|
@ -18348,3 +18934,38 @@ async def test_key_creator_cannot_detach_project_without_admin_access():
|
|||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_runs_custom_key_policy_per_key(monkeypatch):
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
|
||||
mock = _setup_team_keys_mocks(
|
||||
monkeypatch, find_many=keys, update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})})
|
||||
)
|
||||
received: list[CustomKeyPolicyRequest] = []
|
||||
|
||||
async def freeze_tok_b(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
|
||||
received.append(policy_request)
|
||||
if policy_request.existing_key is not None and policy_request.existing_key.token == "tok-b":
|
||||
return {"decision": False, "message": "tok-b is frozen"}
|
||||
return {"decision": True}
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", freeze_tok_b)
|
||||
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc", key_ids=["tok-a", "tok-b"], update_fields=KeyUpdateFields(max_budget=50.0)
|
||||
)
|
||||
)
|
||||
|
||||
assert [update.key for update in response.successful_updates] == ["tok-a"]
|
||||
assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [("tok-b", "tok-b is frozen")]
|
||||
mock.update_data.assert_awaited_once()
|
||||
assert [policy_request.operation for policy_request in received] == ["update", "update"]
|
||||
assert [policy_request.effective_key.max_budget for policy_request in received] == [50.0, 50.0]
|
||||
assert [policy_request.effective_key.team_id for policy_request in received] == ["team-abc", "team-abc"]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
],
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue