[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 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
2026-03-16 21:35:21 -07:00
co-authored by Claude Opus 4.6
parent a771fe55e4
commit 53d96c8353
4 changed files with 176 additions and 6 deletions
@@ -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}
@@ -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
@@ -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):
@@ -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>