diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 87078a5df2..e09d7607eb 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -72,6 +72,9 @@ from litellm.proxy.management_helpers.team_member_permission_checks import ( ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key +from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + get_ui_settings_cached, +) from litellm.proxy.utils import ( PrismaClient, ProxyLogging, @@ -96,6 +99,24 @@ from litellm.types.utils import ( ) +async def _check_custom_key_allowed(custom_key_value: Optional[str]) -> None: + """Raise 403 if custom API keys are disabled and a custom key was provided.""" + if custom_key_value is None: + return + + ui_settings = await get_ui_settings_cached() + if ui_settings.get("disable_custom_api_keys", False) is True: + verbose_proxy_logger.warning( + "Custom API key rejected: disable_custom_api_keys is enabled" + ) + raise HTTPException( + status_code=403, + detail={ + "error": "Custom API key values are disabled by your administrator. Keys must be auto-generated." + }, + ) + + def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]): return data.team_id is not None @@ -671,6 +692,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 prisma_client=prisma_client, ) + # Reject custom key values if disabled by admin + await _check_custom_key_allowed(data.key) + # Validate user-provided key format if data.key is not None and not data.key.startswith("sk-"): _masked = ( @@ -3479,8 +3503,10 @@ async def _rotate_master_key( # noqa: PLR0915 ) -def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: +async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: if data and data.new_key is not None: + # Reject custom key values if disabled by admin + await _check_custom_key_allowed(data.new_key) new_token = data.new_key if not data.new_key.startswith("sk-"): raise HTTPException( @@ -3572,7 +3598,7 @@ async def _execute_virtual_key_regeneration( """Generate new token, update DB, invalidate cache, and return response.""" from litellm.proxy.proxy_server import hash_token - new_token = get_new_token(data=data) + new_token = await get_new_token(data=data) new_token_hash = hash_token(new_token) new_token_key_name = f"sk-...{new_token[-4:]}" update_data = {"token": new_token_hash, "key_name": new_token_key_name} diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 0fa27905ba..d31079f14f 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -149,6 +149,7 @@ ALLOWED_UI_SETTINGS_FIELDS = { "disable_vector_stores_for_internal_users", "allow_vector_stores_for_team_admins", "scope_user_search_to_org", + "disable_custom_api_keys", } # Flags that must be synced from the persisted UISettings into 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 ed2607da24..a3bca77ae3 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 @@ -960,22 +960,34 @@ async def test_key_info_returns_object_permission(monkeypatch): ) -def test_get_new_token_with_valid_key(): +@pytest.mark.asyncio +async def test_get_new_token_with_valid_key(monkeypatch): """Test get_new_token function when provided with a valid key that starts with 'sk-'""" + from unittest.mock import AsyncMock + from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( get_new_token, ) + # Mock get_ui_settings_cached to return setting disabled (custom keys allowed) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + # Test with valid new_key data = RegenerateKeyRequest(new_key="sk-test123456789") - result = get_new_token(data) + result = await get_new_token(data) assert result == "sk-test123456789" -def test_get_new_token_with_invalid_key(): +@pytest.mark.asyncio +async def test_get_new_token_with_invalid_key(monkeypatch): """Test get_new_token function when provided with an invalid key that doesn't start with 'sk-'""" + from unittest.mock import AsyncMock + from fastapi import HTTPException from litellm.proxy._types import RegenerateKeyRequest @@ -983,16 +995,145 @@ def test_get_new_token_with_invalid_key(): get_new_token, ) + # Mock get_ui_settings_cached to return setting disabled (custom keys allowed) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + # Test with invalid new_key (doesn't start with 'sk-') data = RegenerateKeyRequest(new_key="invalid-key-123") with pytest.raises(HTTPException) as exc_info: - get_new_token(data) + await get_new_token(data) assert exc_info.value.status_code == 400 assert "New key must start with 'sk-'" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_check_custom_key_allowed_when_disabled(monkeypatch): + """_check_custom_key_allowed raises 403 when disable_custom_api_keys is true.""" + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + with pytest.raises(HTTPException) as exc_info: + await _check_custom_key_allowed("sk-custom-key-123") + + assert exc_info.value.status_code == 403 + assert "disabled" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_check_custom_key_allowed_when_enabled(monkeypatch): + """_check_custom_key_allowed does nothing when disable_custom_api_keys is false.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": False}), + ) + + # Should not raise + await _check_custom_key_allowed("sk-custom-key-123") + + +@pytest.mark.asyncio +async def test_check_custom_key_allowed_when_unset(monkeypatch): + """_check_custom_key_allowed does nothing when setting is not present.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + # Should not raise + await _check_custom_key_allowed("sk-custom-key-123") + + +@pytest.mark.asyncio +async def test_check_custom_key_allowed_none_key_always_passes(monkeypatch): + """_check_custom_key_allowed does nothing when key is None, even if setting is on.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + # Should not raise — None means auto-generate + await _check_custom_key_allowed(None) + + +@pytest.mark.asyncio +async def test_get_new_token_rejected_when_custom_keys_disabled(monkeypatch): + """get_new_token raises 403 when new_key is set and disable_custom_api_keys is true.""" + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + get_new_token, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + data = RegenerateKeyRequest(new_key="sk-custom-regen-key") + + with pytest.raises(HTTPException) as exc_info: + await get_new_token(data) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_new_token_auto_generates_when_custom_keys_disabled(monkeypatch): + """get_new_token auto-generates a key when new_key is None, even if setting is on.""" + from unittest.mock import AsyncMock + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + get_new_token, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + data = RegenerateKeyRequest() # no new_key + result = await get_new_token(data) + + assert result.startswith("sk-") + + @pytest.mark.asyncio async def test_generate_service_account_requires_team_id(): with pytest.raises(HTTPException): diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 71e882db3f..d0a7a909d6 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -166,6 +166,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const { data: projects, isLoading: isProjectsLoading } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); + const disableCustomApiKeys = Boolean(uiSettingsData?.values?.disable_custom_api_keys); const queryClient = useQueryClient(); const [form] = Form.useForm(); const [isModalVisible, setIsModalVisible] = useState(false); @@ -1581,6 +1582,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp "budget_duration", "tpm_limit", "rpm_limit", + ...(disableCustomApiKeys ? ["key"] : []), ]} />