Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 28 additions & 2 deletions litellm/proxy/management_endpoints/key_management_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,9 @@
)
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,
Expand All @@ -96,6 +99,24 @@
)


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:
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Fragile strict is True equality check

ui_settings.get("disable_custom_api_keys", False) is True will only match the Python boolean True. If the value is stored/deserialized as the string "true" or the integer 1 (both of which are truthy), the guard silently passes and the setting is ignored. The standard Python idiom is a simple truthiness check:

Suggested change
if ui_settings.get("disable_custom_api_keys", False) is True:
if ui_settings.get("disable_custom_api_keys", False):

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."
},
)
Comment on lines +102 to +117
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 DB call on every key-generation request (critical path)

_check_custom_key_allowed calls get_ui_settings_cached() unconditionally — even when custom_key_value is None (i.e., the common case where the user does not supply a custom key). The early-return only fires before get_ui_settings_cached() is invoked, so an unnecessary cache/DB lookup still occurs for every auto-generated key request. Per project rules, critical-path code must avoid unnecessary DB/cache requests.

Move the None guard so that the settings lookup is skipped entirely when no custom key is being requested:

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  # no custom key – nothing to check

    ui_settings = await get_ui_settings_cached()
    if ui_settings.get("disable_custom_api_keys", False):
        ...

Wait — that IS the current code. The issue is that the None check at line 104 returns before the get_ui_settings_cached() call, which is correct. However, _common_key_generation_helper (line 696) and get_new_token (line 3509) call this function inside the request path on every key-generation and key-regeneration call. While the cache mitigates the DB overhead, the async cache lookup still adds latency on every request even when custom_key_value is None would have bailed out first — specifically in get_new_token, the call happens even for the auto-generate branch (data.new_key is None) since _check_custom_key_allowed is called after the if data and data.new_key is not None guard.

Actually, let me re-read:

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)

OK so in get_new_token, the call is inside the if data and data.new_key is not None: guard, so the None check inside _check_custom_key_allowed is redundant but not harmful for that path.

For _common_key_generation_helper (line 696), data.key may be None, and _check_custom_key_allowed is called unconditionally — but then it returns early at the None check on line 104. So actually there is no extra cache hit when data.key is None.

Wait, but that means every call to _common_key_generation_helper — even those where no custom key is provided — invokes an async function that starts executing (just to check None). The function itself does not make a cache call when custom_key_value is None, due to the early return. So actually the implementation is correct on performance grounds.

Let me remove this comment and only keep the stricter logic issue (is True vs truthy).

Actually, let me reconsider the is True issue. This is still a valid concern. Let me check what value JSON deserialization would give - JSON booleans are parsed as Python booleans, so True/False. If the value was stored as an integer by accident it would fail. But for the standard JSON boolean, is True would work correctly. Let me think about whether this is a real issue...

Actually, looking at the code in get_ui_settings_cached:

ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw)

JSON parsing will give proper Python booleans for JSON booleans. So is True would work correctly for properly stored values. However, using is True instead of a truthiness check is still a style concern - it's non-idiomatic and could fail in edge cases (e.g., if someone stores 1 or "true" directly).

I'll keep this as a style/logic note.

Now let me think about other issues:

  1. UI bypass: The UI hides the key field, but this is purely a client-side mitigation. A malicious user could still send the key field in the JSON body directly (bypassing the UI). The backend check in _check_custom_key_allowed is the actual enforcement mechanism, and it does work at the backend level. This is the correct design.

  2. No exemption for proxy admins: The PR description says "No exemptions — proxy admins are also blocked when the setting is enabled." This could be a design concern - maybe proxy admins should be able to bypass this restriction? But it's stated as an explicit design decision, so I'll mention it as a design note rather than a bug.

  3. Cache invalidation: The 10-minute TTL means that if an admin disables disable_custom_api_keys, there's a window where users can still create custom keys. This is acceptable for a UI setting.

  4. Test mock patching: In the tests, the monkeypatch targets litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached. This is correct as the function is imported directly.

  5. Let me think about whether there's anything in the UI changes that has issues...

Looking at the UI:

const disableCustomApiKeys = Boolean(uiSettingsData?.values?.disable_custom_api_keys);

This just hides the field. The backend does the actual enforcement.

  1. Missing UI feedback: When disable_custom_api_keys is enabled, the key field is simply hidden from the "Advanced Settings" accordion. There's no visible message to users explaining why the field is missing. This is a UX improvement suggestion, not a bug.

  2. The is True check: As I discussed, this is using strict identity comparison instead of truthiness. This is a potential logic issue if values are stored in different formats.

Let me focus on the most meaningful issues:

  1. The is True strict comparison (could silently fail to enforce the restriction if value isn't Python True)
  2. The PR description says "No exemptions — proxy admins are also blocked" - this might be a design concern worth flagging, as it deviates from usual patterns where proxy admins get elevated permissions

Let me also check if there's any issue with the UI change - does the form still submit the key field even when hidden? Looking at the code, excludedFields should prevent it. Let me verify by reading more context...

Actually, let me check what excludedFields does in this component - does it prevent the field from being submitted, or just hide it visually? If it only hides the field visually but still submits the value (e.g., if the form had a pre-filled value), there could be an issue where the UI sends the field and the backend rejects it.

Actually I think the excludedFields likely excludes the field from the form schema/display entirely, so it wouldn't be submitted. The backend check is still the authoritative guard.

  1. Stale cache on toggle: If an admin enables disable_custom_api_keys, existing users with the UI still showing the key field (before their browser refreshes) could still submit a custom key. The backend would reject it (after the cache propagates), but the UX would be jarring. This is a minor UX issue.

Let me look at the is True check more carefully. This is the main logic issue worth flagging.

Let me also think about race conditions - get_ui_settings_cached reads from cache or DB. If the admin has just changed the setting, there's a 10-minute window where custom keys might still be accepted. But this is the expected behavior with TTL caching and is not a bug per se.

Let me finalize the review now.



def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]):
return data.team_id is not None

Expand Down Expand Up @@ -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)
Comment on lines +695 to +696
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Fail-fast: move check before alias-uniqueness lookup

The _check_custom_key_allowed call is positioned after _enforce_unique_key_alias, which performs a database lookup whenever key_alias is set. When the disable_custom_api_keys flag is enabled, a request that provides both a custom key and a key_alias will trigger the alias-uniqueness lookup unnecessarily before receiving the 403 response.

Moving _check_custom_key_allowed to before _validate_key_alias_format and _enforce_unique_key_alias ensures the security guard fires before any expensive lookups, which is the standard fail-fast pattern for access controls.


# Validate user-provided key format
if data.key is not None and not data.key.startswith("sk-"):
_masked = (
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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}
Expand Down
6 changes: 6 additions & 0 deletions litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand All @@ -149,6 +154,7 @@ class UISettingsResponse(SettingsResponse):
"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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -960,39 +960,180 @@ 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
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 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,
)
Comment on lines +1060 to +1080
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Unnecessary mock in None-key test

_check_custom_key_allowed(None) returns early at line 104 (if custom_key_value is None: return) before ever calling get_ui_settings_cached. The AsyncMock set up in test_check_custom_key_allowed_none_key_always_passes is therefore never invoked, which means the test doesn't actually verify that the cache is skipped on a None key — and the mock gives a false sense of security about the DB-bypass guarantee.

Consider either removing the mock entirely (which makes the contract explicit: no DB call needed at all), or asserting that the mock was not called to document the expected fast-path behaviour:

mock_get_settings = AsyncMock(return_value={"disable_custom_api_keys": True})
monkeypatch.setattr(
    "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
    mock_get_settings,
)
await _check_custom_key_allowed(None)
mock_get_settings.assert_not_awaited()  # None key must not trigger a cache/DB lookup


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-")
Comment on lines +1117 to +1134
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Unreachable mock in auto-generate test

The AsyncMock for get_ui_settings_cached set up in this test is never invoked. When data.new_key is None, get_new_token enters the else branch immediately:

async def get_new_token(data):
    if data and data.new_key is not None:   # ← False when new_key is None
        await _check_custom_key_allowed(...)  # ← never reached
    else:
        new_token = f"sk-{secrets.token_urlsafe(...)}"  # ← always taken here

_check_custom_key_allowed (and therefore get_ui_settings_cached) is never called, so the mock has no effect on the test outcome. The mock creates a false impression that this test verifies cache-bypass behaviour, similar to the pattern already identified in test_check_custom_key_allowed_none_key_always_passes. Consider removing the mock to make the test contract explicit, or asserting mock_get_settings.assert_not_awaited() to document the expected fast-path.

Suggested change
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_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,
)
mock_get_settings = AsyncMock(return_value={"disable_custom_api_keys": True})
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
mock_get_settings,
)
data = RegenerateKeyRequest() # no new_key
result = await get_new_token(data)
assert result.startswith("sk-")
mock_get_settings.assert_not_awaited() # None new_key must not trigger cache/DB



@pytest.mark.asyncio
async def test_generate_service_account_requires_team_id():
with pytest.raises(HTTPException):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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 ? (
Expand Down Expand Up @@ -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}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -1581,6 +1582,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
"budget_duration",
"tpm_limit",
"rpm_limit",
...(disableCustomApiKeys ? ["key"] : []),
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Stale form value may survive field exclusion

When disableCustomApiKeys transitions from false to true mid-session (e.g. the admin toggles the setting while a user has the "Create Key" modal open), Ant Design's Form instance retains any value the user had already typed into the key field even after it is removed from the rendered fields. The SchemaFormFields useEffect re-runs because excludedFields changed, but it only calls form.setFieldsValue with default values for non-excluded fields — it never clears the now-excluded key field. As a result form.validateFields() will still include key: "sk-..." in the submitted payload, causing a confusing 403 from the backend for a field the user can no longer even see.

The fix is to explicitly clear the field when the flag becomes active:

useEffect(() => {
  if (disableCustomApiKeys) {
    form.setFieldValue("key", undefined);
  }
}, [disableCustomApiKeys, form]);

Add this useEffect alongside the existing form / disableCustomApiKeys declarations.

]}
/>
</AccordionBody>
Expand Down
Loading