mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-22 10:25:58 +00:00
[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:
co-authored by
Claude Opus 4.6
parent
a771fe55e4
commit
53d96c8353
@@ -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>
|
||||
|
||||
Reference in New Issue
Block a user