diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8b402fcff77..95ea8ad4187 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -512,6 +512,14 @@ async def _common_key_generation_helper( # noqa: PLR0915 ) if upperbound_value is not None: if value is None: + if key == "duration" and key in data.model_fields_set: + # Explicitly null = never expires → exceeds any finite upperbound + raise HTTPException( + status_code=400, + detail={ + "error": f"{key} is over max limit set in config - user_value=null (never expires); max_value={upperbound_value}" + }, + ) # Use the upperbound value if user didn't provide a value setattr(data, key, upperbound_value) else: diff --git a/tests/test_litellm/enterprise/proxy/test_key_duration_ceiling.py b/tests/test_litellm/enterprise/proxy/test_key_duration_ceiling.py index 1e6541dbcec..8be95e8cbdc 100644 --- a/tests/test_litellm/enterprise/proxy/test_key_duration_ceiling.py +++ b/tests/test_litellm/enterprise/proxy/test_key_duration_ceiling.py @@ -17,21 +17,27 @@ from unittest.mock import MagicMock from litellm.proxy._types import GenerateKeyRequest, LiteLLM_TeamTable +_ENTERPRISE_SOURCE = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + "..", "..", "..", "..", + "enterprise", "litellm_enterprise", "proxy", + "management_endpoints", "key_management_endpoints.py", + ) +) + +pytestmark = pytest.mark.skipif( + not os.path.exists(_ENTERPRISE_SOURCE), + reason=f"Enterprise source not available at {_ENTERPRISE_SOURCE}", +) + def _load_local_add_team_member_key_duration(): """Load add_team_member_key_duration from the local enterprise source tree.""" - local_path = os.path.normpath( - os.path.join( - os.path.dirname(__file__), - "..", "..", "..", "..", - "enterprise", "litellm_enterprise", "proxy", - "management_endpoints", "key_management_endpoints.py", - ) - ) module_name = "_local_enterprise_key_management_endpoints" # Remove cached version so we always reload from the local file sys.modules.pop(module_name, None) - spec = importlib.util.spec_from_file_location(module_name, local_path) + spec = importlib.util.spec_from_file_location(module_name, _ENTERPRISE_SOURCE) module = importlib.util.module_from_spec(spec) # type: ignore[arg-type] spec.loader.exec_module(module) # type: ignore[union-attr] return module.add_team_member_key_duration @@ -109,10 +115,12 @@ class TestAddTeamMemberKeyDuration: assert result.duration is None def test_service_account_returns_unchanged(self): - """user_id=None (service account) → data is returned unchanged.""" + """user_id=None (service account) → team duration is NOT applied.""" fn = _load_local_add_team_member_key_duration() data = GenerateKeyRequest(user_id=None) result = fn(_make_team("30d"), data) - assert result.duration is None + # Verify the service-account guard prevented the team max from being applied + assert result.duration is None, "Service account should not inherit team duration" + assert result.duration != "30d" 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 a1605734d8e..fc2ecce9df8 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 @@ -6457,6 +6457,75 @@ class TestValidateKeyAliasFormat: assert "Invalid key_alias format" in str(exc.value.message) +class TestUpperboundNullDurationValidation: + """Tests for upperbound validation when duration is explicitly null (never expires).""" + + @pytest.mark.asyncio + async def test_explicit_null_duration_raises_when_upperbound_set(self): + """duration=null (never expires) should raise 400 when upperbound duration is set.""" + import litellm + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = ( + LiteLLM_UpperboundKeyGenerateParams(duration="30d") + ) + with pytest.raises(HTTPException) as exc_info: + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1", duration=None), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + assert exc_info.value.status_code == 400 + assert "never expires" in str(exc_info.value.detail).lower() + finally: + litellm.upperbound_key_generate_params = original + + @pytest.mark.asyncio + async def test_omitted_duration_uses_upperbound_as_default(self): + """Duration not sent → upperbound is applied as default (no error).""" + import litellm + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = ( + LiteLLM_UpperboundKeyGenerateParams(duration="30d") + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={ + "key": "sk-test", + "expires": None, + "user_id": "user-1", + }, + ), patch("litellm.proxy.proxy_server.prisma_client"), patch( + "litellm.proxy.proxy_server.llm_router" + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ): + # Should not raise — duration omitted, so upperbound fills in as default + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + finally: + litellm.upperbound_key_generate_params = original + + class TestCommonKeyGenerationHelperTeamDurationValidation: """Tests for team duration validation inside _common_key_generation_helper."""