From 53d96c8353b2fd18669a54ec0ff2410ff3518ad5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 16 Mar 2026 21:35:21 -0700 Subject: [PATCH 1/5] [Feature] Disable custom API key values via UI setting Add disable_custom_api_keys UI setting that prevents users from specifying custom key values during key generation and regeneration. When enabled, all keys must be auto-generated, eliminating the risk of key hash collisions in multi-tenant environments. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 30 +++- .../proxy_setting_endpoints.py | 1 + .../test_key_management_endpoints.py | 149 +++++++++++++++++- .../organisms/create_key_button.tsx | 2 + 4 files changed, 176 insertions(+), 6 deletions(-) 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"] : []), ]} /> From 72aa5fc21926865d3aa6b20afad11c51273539ac Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 16 Mar 2026 22:10:02 -0700 Subject: [PATCH 2/5] [Fix] Add disable_custom_api_keys to UISettings Pydantic model Without this field on the model, GET /get/ui_settings omits the setting from the response and field_schema, preventing the UI from reading it. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index d31079f14f..9faadd334a 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -129,6 +129,11 @@ class UISettings(BaseModel): description="If enabled, the user search endpoint (/user/filter/ui) restricts results by organization. When off, any authenticated user can search all users.", ) + disable_custom_api_keys: bool = Field( + default=False, + description="If true, users cannot specify custom API key values. All keys must be auto-generated.", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" From c687e631b4581afcf9df057f9ad8e5d9c6ec7aa1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 16 Mar 2026 22:26:11 -0700 Subject: [PATCH 3/5] [Feature] Add disable_custom_api_keys toggle to UI Settings page Adds a toggle switch to the admin UI Settings page so administrators can enable/disable custom API key values without making direct API calls. Co-Authored-By: Claude Opus 4.6 --- .../AdminSettings/UISettings/UISettings.tsx | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index a22c78c943..0d12a4c4bf 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -24,6 +24,7 @@ export default function UISettings() { const disableVectorStoresProperty = schema?.properties?.disable_vector_stores_for_internal_users; const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org; + const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); @@ -182,6 +183,20 @@ export default function UISettings() { ); }; + const handleToggleDisableCustomApiKeys = (checked: boolean) => { + updateSettings( + { disable_custom_api_keys: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + return ( {isLoading ? ( @@ -382,6 +397,26 @@ export default function UISettings() { + {/* Disable custom API key values */} + + + + Disable custom API key values + + {disableCustomApiKeysProperty?.description ?? + "If true, users cannot specify custom API key values. All keys must be auto-generated."} + + + + + + {/* Page Visibility for Internal Users */} Date: Mon, 16 Mar 2026 22:26:53 -0700 Subject: [PATCH 4/5] [Fix] Rename toggle label to "Disable custom Virtual key values" Co-Authored-By: Claude Opus 4.6 --- .../Settings/AdminSettings/UISettings/UISettings.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 0d12a4c4bf..f3eb6abdec 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -397,17 +397,17 @@ export default function UISettings() { - {/* Disable custom API key values */} + {/* Disable custom Virtual key values */} - Disable custom API key values + Disable custom Virtual key values {disableCustomApiKeysProperty?.description ?? "If true, users cannot specify custom API key values. All keys must be auto-generated."} From 471e0f147e3a19da4c55e7b54a92c4023a07f49c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 16 Mar 2026 22:35:44 -0700 Subject: [PATCH 5/5] [Fix] Remove "API" from custom key description text Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py | 2 +- .../components/Settings/AdminSettings/UISettings/UISettings.tsx | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 9faadd334a..60bf41709e 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -131,7 +131,7 @@ class UISettings(BaseModel): disable_custom_api_keys: bool = Field( default=False, - description="If true, users cannot specify custom API key values. All keys must be auto-generated.", + description="If true, users cannot specify custom key values. All keys must be auto-generated.", ) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index f3eb6abdec..fb7c38449b 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -410,7 +410,7 @@ export default function UISettings() { Disable custom Virtual key values {disableCustomApiKeysProperty?.description ?? - "If true, users cannot specify custom API key values. All keys must be auto-generated."} + "If true, users cannot specify custom key values. All keys must be auto-generated."}