scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

feat: add allow_custom_api_keys setting to block user-specified API keys

Browse source 
- Add `allow_custom_api_keys` to ConfigGeneralSettings and UISettings
- Block custom keys in _common_key_generation_helper and get_new_token when disabled
- Hide custom key field in UI when setting is false
- Fix: only exempt PROXY_ADMIN (not VIEW_ONLY) from user_id mutation guards
- Add test assertion for allowed custom key path
This commit is contained in:
Ryan Crabbe 2026-03-09 16:28:54 -07:00
parent a23ab1972c
commit e73a50f9a3
5 changed files with 126 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

@ -2203,6 +2203,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.",
)
allow_custom_api_keys: Optional[bool] = Field(
None,
description="If False, users cannot specify a custom API key value when creating keys. Keys will always be randomly generated.",
)
class ConfigYAML(LiteLLMPydanticObjectBase):

33
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -650,6 +650,18 @@ async def _common_key_generation_helper( # noqa: PLR0915
prisma_client=prisma_client,
)
# Block custom API keys when disabled in settings
if data.key is not None:
from litellm.proxy.proxy_server import general_settings
if not general_settings.get("allow_custom_api_keys", True):
raise HTTPException(
status_code=403,
detail={
"error": "Custom API keys are disabled. Keys must be randomly generated."
},
)
# Validate user-provided key format
if data.key is not None and not data.key.startswith("sk-"):
_masked = (
@ -1168,10 +1180,9 @@ async def generate_key_fn(
# Auto-populate user_id from authenticated user when not provided (non-admins only)
if data.user_id is None and user_api_key_dict.user_id is not None:
if user_api_key_dict.user_role not in [
LitellmUserRoles.PROXY_ADMIN.value,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
]:
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
):
data.user_id = user_api_key_dict.user_id
# Validate budget values are not negative
@ -1894,10 +1905,7 @@ async def update_key_fn(
and _update_fields["user_id"] is None
and existing_key_row.user_id is not None
and user_api_key_dict.user_role
not in [
LitellmUserRoles.PROXY_ADMIN.value,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
]
!= LitellmUserRoles.PROXY_ADMIN.value
):
raise HTTPException(
status_code=403,
@ -3333,6 +3341,15 @@ async def _rotate_master_key( # noqa: PLR0915
def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
if data and data.new_key is not None:
from litellm.proxy.proxy_server import general_settings
if not general_settings.get("allow_custom_api_keys", True):
raise HTTPException(
status_code=403,
detail={
"error": "Custom API keys are disabled. Keys must be randomly generated."
},
)
new_token = data.new_key
if not data.new_key.startswith("sk-"):
raise HTTPException(

7
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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.",
)
allow_custom_api_keys: bool = Field(
default=True,
description="If false, users cannot specify a custom API key value when creating keys. Keys will always be randomly 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",
"allow_custom_api_keys",
}
# Flags that must be synced from the persisted UISettings into
@ -159,6 +165,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS = [
"allow_agents_for_team_admins",
"disable_vector_stores_for_internal_users",
"allow_vector_stores_for_team_admins",
"allow_custom_api_keys",
]

88
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -986,6 +986,25 @@ def test_get_new_token_with_invalid_key():
assert "New key must start with 'sk-'" in str(exc_info.value.detail)
def test_get_new_token_rejected_when_custom_keys_disabled():
"""Test get_new_token rejects custom keys when allow_custom_api_keys is False"""
from litellm.proxy._types import RegenerateKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
get_new_token,
)
data = RegenerateKeyRequest(new_key="sk-custom-regen-key")
with patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_custom_api_keys": False},
):
with pytest.raises(HTTPException) as exc_info:
get_new_token(data)
assert exc_info.value.status_code == 403
assert "Custom API keys are disabled" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_generate_service_account_requires_team_id():
with pytest.raises(HTTPException):
@ -6865,3 +6884,72 @@ async def test_generate_key_validation_user_id_and_team_id(monkeypatch):
),
)
assert exc.value.code == "403"
@pytest.mark.asyncio
async def test_custom_api_key_rejected_when_disabled():
"""
When general_settings['allow_custom_api_keys'] is False,
_common_key_generation_helper should reject requests that specify a custom key.
"""
with patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_custom_api_keys": False},
), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.llm_router"
), patch(
"litellm.proxy.proxy_server.premium_user", False
):
mock_prisma.return_value = AsyncMock()
with pytest.raises(HTTPException) as exc:
await _common_key_generation_helper(
data=GenerateKeyRequest(key="sk-custom-key-123"),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234"
),
litellm_changed_by=None,
team_table=None,
)
assert exc.value.status_code == 403
assert "Custom API keys are disabled" in exc.value.detail["error"]
@pytest.mark.asyncio
async def test_custom_api_key_allowed_when_enabled():
"""
When general_settings['allow_custom_api_keys'] is True (default),
_common_key_generation_helper should allow custom keys.
"""
with patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_custom_api_keys": True},
), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.llm_router"
), patch(
"litellm.proxy.proxy_server.premium_user", False
), patch(
"litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
) as mock_generate_key:
mock_prisma.return_value = AsyncMock()
mock_generate_key.return_value = {
"key": "sk-custom-key-123",
"expires": None,
"user_id": "test-user",
"team_id": None,
}
# Should NOT raise
await _common_key_generation_helper(
data=GenerateKeyRequest(key="sk-custom-key-123"),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="test-user",
),
litellm_changed_by=None,
team_table=None,
)
mock_generate_key.assert_awaited_once()

2
ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
View file

@ -163,6 +163,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 allowCustomApiKeys = uiSettingsData?.values?.allow_custom_api_keys !== false;
const queryClient = useQueryClient();
const [form] = Form.useForm();
const [isModalVisible, setIsModalVisible] = useState(false);
@ -1537,6 +1538,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
"budget_duration",
"tpm_limit",
"rpm_limit",
...(!allowCustomApiKeys ? ["key"] : []),
]}
/>
</AccordionBody>