From 7d0a906a93c08a32c261a4fb520d127cd0949b62 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 26 Sep 2026 10:38:02 -0500 Subject: [PATCH 1/3] fix(proxy): persist explicit virtual key expiration --- litellm/proxy/_types.py | 1 + .../key_management_endpoints.py | 18 +++++----- .../test_key_management_endpoints.py | 34 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 ++++ 4 files changed, 51 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4597872d84e..cde1c6ce8ba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1269,6 +1269,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 7e159ec90e7..bbcd54c987c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -419,7 +419,9 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite 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 + 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 @@ -4616,6 +4618,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 @@ -4635,12 +4638,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 @@ -4717,7 +4719,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, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index aa6be328f4a..a5729e394c4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13907,6 +13907,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", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d46af04730..98064894a80 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30643,6 +30643,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 */ @@ -39209,6 +39211,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 */ @@ -44591,6 +44595,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 */ From 4956fc8404bbeef3d35f595ead58ee72def54841 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 26 Sep 2026 10:52:29 -0500 Subject: [PATCH 2/3] docs(proxy): document explicit key expiration --- litellm/proxy/management_endpoints/key_management_endpoints.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index bbcd54c987c..6e963b095af 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3410,6 +3410,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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 98064894a80..3e583e2961d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8404,6 +8404,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 From acdccee0fe72c38ecc3d4b156086a541446e3854 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 26 Sep 2026 11:51:50 -0500 Subject: [PATCH 3/3] fix(proxy): enforce lifetime limits for explicit expiration --- .../key_management_endpoints.py | 80 ++++++++----- .../test_key_management_endpoints.py | 111 +++++++++++++++++- 2 files changed, 157 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6e963b095af..db9eb38732b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1109,6 +1109,20 @@ _BUDGET_NUMERIC_KEYS = frozenset( ) +def _key_expiration_exceeds_limit( + data: GenerateKeyRequest | UpdateKeyRequest, maximum_duration: str, fill_defaults: bool +) -> 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 = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(maximum_duration)) + return expires > maximum_expiration + + def _enforce_upperbound_key_params( data: GenerateKeyRequest | UpdateKeyRequest, fill_defaults: bool = True, @@ -1135,40 +1149,41 @@ 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): + 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( @@ -1242,6 +1257,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", diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a5729e394c4..26d37c5e0e2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -85,6 +85,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) @@ -12722,6 +12723,87 @@ 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")) + future: Final = datetime.now(timezone.utc) + 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) + + 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_preserves_bounded_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")) + expires: Final = datetime.now(timezone.utc) + timedelta(hours=1) + data: Final = request_type(key="sk-expiration-test", expires=expires) + + _enforce_upperbound_key_params(data, fill_defaults=fill_defaults) + + 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")) + expires: Final = datetime.now(timezone.utc) + 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) + + 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")) + expires: Final = datetime.now(timezone.utc) + timedelta(days=2) + data: Final = request_type(key="sk-expiration-test", duration="1h", expires=expires) + + _enforce_upperbound_key_params(data, fill_defaults=fill_defaults) + + 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 @@ -12783,7 +12865,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 ( @@ -12801,7 +12886,11 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk ), ) existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest(duration="2h") + data: Final = ( + RegenerateKeyRequest(expires=datetime.now(timezone.utc) + 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() @@ -19310,6 +19399,24 @@ 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) + expires: Final = datetime.now(timezone.utc) + timedelta(hours=1) + + key_row: Final = await _generate_key_and_get_persisted_row(GenerateKeyRequest(expires=expires), insert_data) + + 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,