diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 9783f44b41f..06fcffc00f9 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -339,17 +339,15 @@ async def _sign_in( response = await generate_key_helper_fn( llm_router=None, request_type="key", - **{ - "user_role": LitellmUserRoles.PROXY_ADMIN, - "duration": LITELLM_UI_SESSION_DURATION, - "key_max_budget": litellm.max_ui_session_budget, - "models": [], - "aliases": {}, - "config": {}, - "spend": 0, - "user_id": key_user_id, - "team_id": "litellm-dashboard", - }, + user_role=LitellmUserRoles.PROXY_ADMIN, + duration=LITELLM_UI_SESSION_DURATION, + key_max_budget=litellm.max_ui_session_budget, + models=[], + aliases={}, + config={}, + spend=0, + user_id=key_user_id, + team_id="litellm-dashboard", ) else: raise ProxyException( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index b563dbc5c4c..9cd724c4a12 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -594,7 +594,7 @@ async def new_user( teams = check_if_default_team_set() organization_ids: Final = cast(list[str] | None, data_json.pop("organizations", None)) - response: Final = await generate_key_helper_fn(request_type="user", **data_json, llm_router=None) + response: Final = await generate_key_helper_fn(request_type="user", expires=None, **data_json, llm_router=None) created_user_id: Final = cast(str | None, response.get("user_id", None)) if created_user_id is not None: forget_missing_user(created_user_id) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 824b310a1dd..4fd4fb306ac 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -479,7 +479,9 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite else {**preset_columns, "allowed_routes": list(default_allowed_routes)} ) 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 @@ -1318,9 +1320,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. @@ -1344,40 +1362,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}" + }, + ) def _team_admin_requests_member_key( @@ -1416,6 +1437,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 ( @@ -1501,6 +1524,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", @@ -1522,7 +1546,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) _enforce_budget_delegation_ceiling( requested_max_budget=_requested_max_budget, @@ -1780,7 +1804,11 @@ async def _common_key_generation_helper( try: response = await generate_key_helper_fn( - request_type="key", **data_json, table_name="key", llm_router=llm_router + request_type="key", + expires=data.expires, + **{key: value for key, value in data_json.items() if key != "expires"}, + table_name="key", + llm_router=llm_router, ) except KeyProjectBindingError: if prisma_client is not None: @@ -3856,6 +3884,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 @@ -5086,6 +5115,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 @@ -5105,12 +5135,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 @@ -5187,7 +5216,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, @@ -6046,6 +6075,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 @@ -6107,7 +6137,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/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 78187a59e5b..8257911a9da 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18827,17 +18827,15 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str: response: Final = await generate_key_helper_fn( llm_router=llm_router, request_type="key", - **{ - "user_role": user_obj.user_role, - "duration": LITELLM_UI_SESSION_DURATION, - "key_max_budget": litellm.max_ui_session_budget, - "models": [], - "aliases": {}, - "config": {}, - "spend": 0, - "user_id": user_obj.user_id, - "team_id": UI_TEAM_ID, - }, + user_role=user_obj.user_role, + duration=LITELLM_UI_SESSION_DURATION, + key_max_budget=litellm.max_ui_session_budget, + models=[], + aliases={}, + config={}, + spend=0, + user_id=user_obj.user_id, + team_id=UI_TEAM_ID, ) key: Final = response["token"] diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index baa02b68d98..ec0373af974 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -19,6 +19,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/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index d6ee9b1d29e..6b223187a35 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -15,7 +15,7 @@ from fastapi.testclient import TestClient from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter import inspect @@ -104,6 +104,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) @@ -13010,6 +13011,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 @@ -13071,7 +13159,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 ( @@ -13089,7 +13180,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() @@ -13123,6 +13219,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) @@ -14200,6 +14297,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", @@ -19624,7 +19755,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( @@ -19634,9 +19767,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 @@ -19946,6 +20080,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 f3821a2bb46..574430b45d1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8925,6 +8925,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 @@ -32812,6 +32813,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 */ @@ -42094,6 +42097,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 */ @@ -48205,6 +48210,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 */