From 7d0a906a93c08a32c261a4fb520d127cd0949b62 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 26 Sep 2026 10:38:02 -0500 Subject: [PATCH] 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 */