diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d46dd3ab3e9..a71881739ba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1360,6 +1360,7 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): class KeyRequestBase(GenerateRequestBase): key: str | None = None + expires: datetime | None = None tpd_limit: int | None = None default_estimated_output_tokens: PositiveInt | None = None default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ac9dfeb136d..fa7e7d9676a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -440,7 +440,9 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite folded_metadata: Final = {**metadata, **metadata_fields} columns: Final = handle_key_type(data, {**column_fields}) expires: Final = ( - now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None + now + timedelta(seconds=duration_in_seconds(duration=data.duration)) + if data.duration is not None + else data.expires ) budget_reset_at: Final = ( get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None @@ -1137,9 +1139,25 @@ _BUDGET_NUMERIC_KEYS = frozenset( ) +def _key_expiration_exceeds_limit( + data: GenerateKeyRequest | UpdateKeyRequest, maximum_duration: str, fill_defaults: bool, now: datetime +) -> bool: + if data.duration: + return False + if not fill_defaults and data.duration is None and "duration" in data.model_fields_set: + return True + if data.expires is None: + return not fill_defaults and "expires" in data.model_fields_set + expires: Final = data.expires if data.expires.tzinfo is not None else data.expires.replace(tzinfo=timezone.utc) + maximum_expiration: Final = now + timedelta(seconds=duration_in_seconds(maximum_duration)) + return expires > maximum_expiration + + def _enforce_upperbound_key_params( data: GenerateKeyRequest | UpdateKeyRequest, fill_defaults: bool = True, + *, + now: datetime | None = None, ) -> None: """ Enforce upperbound limits on key parameters. @@ -1163,40 +1181,43 @@ def _enforce_upperbound_key_params( if litellm.upperbound_key_generate_params is None: return + maximum_duration: Final = litellm.upperbound_key_generate_params.duration + if maximum_duration is not None and _key_expiration_exceeds_limit( + data, maximum_duration, fill_defaults, now if now is not None else datetime.now(timezone.utc) + ): + raise HTTPException( + status_code=400, + detail={"error": f"expires is over max limit set in config - max_duration={maximum_duration}"}, + ) + for elem in data: key, value = elem - upperbound_value = getattr(litellm.upperbound_key_generate_params, key, None) - if upperbound_value is not None: - if value is None: - if fill_defaults: - setattr(data, key, upperbound_value) - else: - if key in [ - "max_budget", - "max_parallel_requests", - "tpm_limit", - "rpm_limit", - ]: - if value > upperbound_value: - raise HTTPException( - status_code=400, - detail={ - "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" - }, - ) - elif key in ["budget_duration", "duration"]: - upperbound_duration = duration_in_seconds(duration=upperbound_value) - if value == "-1": - user_duration = float("inf") - else: - user_duration = duration_in_seconds(duration=value) - if user_duration > upperbound_duration: - raise HTTPException( - status_code=400, - detail={ - "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" - }, - ) + if (upperbound_value := getattr(litellm.upperbound_key_generate_params, key, None)) is None: + continue + if value is None: + if fill_defaults and not (key == "duration" and data.expires is not None): + setattr(data, key, upperbound_value) + continue + if key in ["max_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"]: + if value > upperbound_value: + raise HTTPException( + status_code=400, + detail={ + "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" + }, + ) + elif key in ["budget_duration", "duration"]: + upperbound_duration, user_duration = ( + duration_in_seconds(duration=upperbound_value), + float("inf") if value == "-1" else duration_in_seconds(duration=value), + ) + if user_duration > upperbound_duration: + raise HTTPException( + status_code=400, + detail={ + "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" + }, + ) async def _common_key_generation_helper( @@ -1204,6 +1225,8 @@ async def _common_key_generation_helper( user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None, team_table: LiteLLM_TeamTableCachedObj | None, + *, + now: datetime | None = None, ) -> GenerateKeyResponse: from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( @@ -1270,6 +1293,7 @@ async def _common_key_generation_helper( if ( value is None and (key != "budget_duration" or key not in data.model_fields_set) + and (key != "duration" or data.expires is None) and key in [ "max_budget", @@ -1291,7 +1315,7 @@ async def _common_key_generation_helper( setattr(data, key, litellm.default_key_generate_params.get(key, {})) # check if user set upperbound key/generate params on config.yaml - _enforce_upperbound_key_params(data, fill_defaults=True) + _enforce_upperbound_key_params(data, fill_defaults=True, now=now) # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. @@ -3444,6 +3468,7 @@ async def update_key_fn( - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - allowed_cache_controls: Optional[list] - List of allowed cache control values - duration: Optional[str] - Key validity duration ("30d", "1h", etc.), null to never expire, or "-1" to never expire (deprecated, use null) + - expires: Optional[datetime] - Absolute expiration timestamp. An explicitly supplied duration takes precedence - permissions: Optional[dict] - Key-specific permissions - send_invite_email: Optional[bool] - Send invite email to user_id - guardrails: Optional[List[str]] - List of active guardrails for the key @@ -4648,6 +4673,7 @@ async def generate_key_helper_fn( budget_limits: list | None = None, # multiple concurrent budget windows *, llm_router: Router | None = None, + expires: datetime | None = None, ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -4667,12 +4693,11 @@ async def generate_key_helper_fn( else: token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" - if duration is None: # allow tokens that never expire - expires = None - else: - # Add duration to current time for exact expiration (not standardized reset time) - duration_seconds: Final = duration_in_seconds(duration) - expires = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds) + resolved_expires: Final = ( + datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration)) + if duration is not None + else expires + ) if key_budget_duration is None: # one-time budget key_reset_at = None @@ -4749,7 +4774,7 @@ async def generate_key_helper_fn( key_data: Final = { "token": token, "key_alias": key_alias, - "expires": expires, + "expires": resolved_expires, "models": models, "aliases": aliases_json, "config": config_json, @@ -5601,6 +5626,7 @@ async def _execute_virtual_key_regeneration( litellm_changed_by: str | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + now: datetime | None = None, ) -> GenerateKeyResponse: """Generate new token, update DB, invalidate cache, and return response.""" from litellm.proxy import proxy_server @@ -5663,7 +5689,7 @@ async def _execute_virtual_key_regeneration( 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) + _enforce_upperbound_key_params(data, fill_defaults=False, now=now) non_default_values = await prepare_key_update_data( data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router ) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 0f21f3ee3e0..3dcd0a77c3f 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -13,6 +13,7 @@ from fastapi.testclient import TestClient from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException +from pydantic import TypeAdapter import inspect @@ -85,6 +86,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyResponse, CustomKeyPolicyRequest, ) +from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams client = TestClient(app) @@ -12717,6 +12719,93 @@ def test_enforce_upperbound_duration_over_limit(monkeypatch): assert "duration" in str(exc_info.value.detail) +@pytest.mark.parametrize( + ("request_type", "fill_defaults"), + ((GenerateKeyRequest, True), (UpdateKeyRequest, False), (RegenerateKeyRequest, False)), +) +@pytest.mark.parametrize("with_timezone", (False, True)) +def test_enforce_upperbound_rejects_expiration_beyond_limit( + monkeypatch: pytest.MonkeyPatch, + request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest], + fill_defaults: bool, + with_timezone: bool, +) -> None: + monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d")) + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + future: Final = now + timedelta(days=2) + expires: Final = future if with_timezone else future.replace(tzinfo=None) + data: Final = request_type(key="sk-expiration-test", expires=expires) + + with pytest.raises(HTTPException) as error: + _enforce_upperbound_key_params(data, fill_defaults=fill_defaults, now=now) + + assert error.value.status_code == 400 + assert "expires" in str(error.value.detail) + + +@pytest.mark.parametrize( + ("request_type", "fill_defaults"), + ((GenerateKeyRequest, True), (UpdateKeyRequest, False), (RegenerateKeyRequest, False)), +) +@pytest.mark.parametrize("hours", (1, 24)) +def test_enforce_upperbound_preserves_bounded_expiration( + monkeypatch: pytest.MonkeyPatch, + request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest], + fill_defaults: bool, + hours: int, +) -> None: + monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d")) + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + expires: Final = now + timedelta(hours=hours) + data: Final = request_type(key="sk-expiration-test", expires=expires) + + _enforce_upperbound_key_params(data, fill_defaults=fill_defaults, now=now) + + assert (data.duration, data.expires) == (None, expires) + + +@pytest.mark.parametrize("request_type", (UpdateKeyRequest, RegenerateKeyRequest)) +@pytest.mark.parametrize("clear_duration", (False, True)) +def test_enforce_upperbound_rejects_clearing_expiration( + monkeypatch: pytest.MonkeyPatch, + request_type: type[UpdateKeyRequest] | type[RegenerateKeyRequest], + clear_duration: bool, +) -> None: + monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d")) + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + expires: Final = now + timedelta(hours=1) + data: Final = request_type.model_validate( + {"key": "sk-expiration-test", "duration": None, "expires": expires} + if clear_duration + else {"key": "sk-expiration-test", "expires": None} + ) + + with pytest.raises(HTTPException) as error: + _enforce_upperbound_key_params(data, fill_defaults=False, now=now) + + assert error.value.status_code == 400 + assert "expires" in str(error.value.detail) + + +@pytest.mark.parametrize( + ("request_type", "fill_defaults"), + ((GenerateKeyRequest, True), (UpdateKeyRequest, False), (RegenerateKeyRequest, False)), +) +def test_enforce_upperbound_uses_explicit_duration_over_expiration( + monkeypatch: pytest.MonkeyPatch, + request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest], + fill_defaults: bool, +) -> None: + monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d")) + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + expires: Final = now + timedelta(days=2) + data: Final = request_type(key="sk-expiration-test", duration="1h", expires=expires) + + _enforce_upperbound_key_params(data, fill_defaults=fill_defaults, now=now) + + assert (data.duration, data.expires) == ("1h", expires) + + def test_enforce_upperbound_no_config_is_noop(monkeypatch): """Test that no enforcement happens when upperbound params are not configured.""" import litellm @@ -12778,7 +12867,10 @@ def _make_regenerate_existing_key(): @pytest.mark.asyncio -async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monkeypatch): +@pytest.mark.parametrize("use_absolute_expiration", (False, True)) +async def test_execute_virtual_key_regeneration_rejects_over_limit_duration( + monkeypatch: pytest.MonkeyPatch, use_absolute_expiration: bool +) -> None: """Regenerate must reject durations exceeding upperbound_key_generate_params.duration.""" from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -12796,7 +12888,12 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk ), ) existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest(duration="2h") + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + data: Final = ( + RegenerateKeyRequest(expires=now + timedelta(hours=2)) + if use_absolute_expiration + else RegenerateKeyRequest(duration="2h") + ) user_api_key_dict = _make_regenerate_user_api_key_dict() mock_prisma_client = _make_regenerate_mock_prisma() @@ -12830,6 +12927,7 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk litellm_changed_by=None, user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), + now=now, ) assert exc_info.value.status_code == 400 assert "duration" in str(exc_info.value.detail) @@ -13902,6 +14000,40 @@ def test_effective_key_for_generate_without_duration_never_expires(): assert effective_key.key_type == "default" +@pytest.mark.parametrize("duration", (None, "1h")) +def test_effective_key_for_generate_preserves_explicit_expiration(duration: str | None) -> None: + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + explicit: Final = now + timedelta(days=1) + request: Final = GenerateKeyRequest(expires=explicit, duration=duration) + + effective: Final = _effective_key_for_generate(data=request, now=now) + + assert effective.expires == (explicit if duration is None else now + timedelta(hours=1)) + assert request.expires == explicit + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_preserves_explicit_expiration() -> None: + explicit: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + existing: Final = LiteLLM_VerificationToken(token="test-token", metadata={}) + request: Final = UpdateKeyRequest(key="test-token", expires=explicit) + + update: Final = await prepare_key_update_data(data=request, existing_key_row=existing) + + assert update["expires"] == explicit + + +@pytest.mark.asyncio +async def test_generate_key_helper_persists_explicit_expiration(monkeypatch: pytest.MonkeyPatch) -> None: + explicit: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + database: Final = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", database) + + await generate_key_helper_fn(request_type="key", expires=explicit, table_name="key", token="sk-expiration-test-token") + + assert database.insert_data.call_args.kwargs["data"]["expires"] == explicit + + def _policy_request_for_generate() -> CustomKeyPolicyRequest: return CustomKeyPolicyRequest( operation="generate", @@ -19323,7 +19455,9 @@ def _wire_key_generation_prisma(monkeypatch): return mock_prisma_client.insert_data -async def _generate_key_and_get_persisted_row(data: GenerateKeyRequest, mock_insert_data): +async def _generate_key_and_get_persisted_row( + data: GenerateKeyRequest, mock_insert_data: AsyncMock, *, now: datetime | None = None +) -> Mapping[str, object]: await _common_key_generation_helper( data=data, user_api_key_dict=UserAPIKeyAuth( @@ -19333,9 +19467,10 @@ async def _generate_key_and_get_persisted_row(data: GenerateKeyRequest, mock_ins ), litellm_changed_by=None, team_table=None, + now=now, ) - key_call = next(c for c in mock_insert_data.call_args_list if c.kwargs["table_name"] == "key") - return key_call.kwargs["data"] + key_call: Final = next(c for c in mock_insert_data.call_args_list if c.kwargs["table_name"] == "key") + return TypeAdapter(Mapping[str, object]).validate_python(key_call.kwargs["data"]) @pytest.mark.asyncio @@ -19366,6 +19501,25 @@ async def test_key_generate_omitted_budget_duration_still_takes_default_key_gene assert key_row["budget_reset_at"] is not None +@pytest.mark.parametrize("with_lifetime_limit", (False, True)) +async def test_key_generate_expiration_overrides_default_duration( + monkeypatch: pytest.MonkeyPatch, with_lifetime_limit: bool +) -> None: + monkeypatch.setattr(litellm, "default_key_generate_params", {"duration": "1d"}) + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams(duration="1d") if with_lifetime_limit else None, + ) + insert_data: Final = _wire_key_generation_prisma(monkeypatch) + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + expires: Final = now + timedelta(hours=1) + + key_row: Final = await _generate_key_and_get_persisted_row(GenerateKeyRequest(expires=expires), insert_data, now=now) + + assert key_row["expires"] == expires + + @pytest.mark.asyncio async def test_key_generate_explicit_null_budget_duration_cannot_bypass_upperbound(monkeypatch): """upperbound_key_generate_params is an admin ceiling: an explicit null must not mint an uncapped key, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fa90213a486..b4e79efeed5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8727,6 +8727,7 @@ export interface paths { * - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" * - allowed_cache_controls: Optional[list] - List of allowed cache control values * - duration: Optional[str] - Key validity duration ("30d", "1h", etc.), null to never expire, or "-1" to never expire (deprecated, use null) + * - expires: Optional[datetime] - Absolute expiration timestamp. An explicitly supplied duration takes precedence * - permissions: Optional[dict] - Key-specific permissions * - send_invite_email: Optional[bool] - Send invite email to user_id * - guardrails: Optional[List[str]] - List of active guardrails for the key @@ -33180,6 +33181,8 @@ export interface components { end_user_budget_id?: string | null; /** Enforced Params */ enforced_params?: string[] | null; + /** Expires */ + expires?: string | null; /** Guardrails */ guardrails?: string[] | null; /** Key */ @@ -42599,6 +42602,8 @@ export interface components { end_user_budget_id?: string | null; /** Enforced Params */ enforced_params?: string[] | null; + /** Expires */ + expires?: string | null; /** Grace Period */ grace_period?: string | null; /** Guardrails */ @@ -48883,6 +48888,8 @@ export interface components { end_user_budget_id?: string | null; /** Enforced Params */ enforced_params?: string[] | null; + /** Expires */ + expires?: string | null; /** Guardrails */ guardrails?: string[] | null; /** Key */