scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 225

Merge 71dacc3f5f into 63a4f3f2f3

Browse source
This commit is contained in:
Emerson Gomes 2026-10-09 07:58:29 -04:00 • committed by GitHub
commit 61467c9d59
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 235 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -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

110
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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
)

164
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -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,

7
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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 */