Merge pull request #23812 from BerriAI/litellm_disable_custom_api_keys

[Feature] Disable Custom Virtual Key Values via UI Setting

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}

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 key values. All keys must be auto-generated.",
+    )
+
 
 class UISettingsResponse(SettingsResponse):
     """Response model for UI settings"""
@@ -149,6 +154,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

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

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 (
     <Card title="UI Settings">
       {isLoading ? (
@@ -382,6 +397,26 @@ export default function UISettings() {
 
           <Divider />
 
+          {/* Disable custom Virtual key values */}
+          <Space align="start" size="middle">
+            <Switch
+              checked={Boolean(values.disable_custom_api_keys)}
+              disabled={isUpdating}
+              loading={isUpdating}
+              onChange={handleToggleDisableCustomApiKeys}
+              aria-label={disableCustomApiKeysProperty?.description ?? "Disable custom Virtual key values"}
+            />
+            <Space direction="vertical" size={4}>
+              <Typography.Text strong>Disable custom Virtual key values</Typography.Text>
+              <Typography.Text type="secondary">
+                {disableCustomApiKeysProperty?.description ??
+                  "If true, users cannot specify custom key values. All keys must be auto-generated."}
+              </Typography.Text>
+            </Space>
+          </Space>
+
+          <Divider />
+
           {/* Page Visibility for Internal Users */}
           <PageVisibilitySettings
             enabledPagesInternalUsers={values.enabled_ui_pages_internal_users}

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -166,6 +166,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ 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<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
                           "budget_duration",
                           "tpm_limit",
                           "rpm_limit",
+                          ...(disableCustomApiKeys ? ["key"] : []),
                         ]}
                       />
                     </AccordionBody>