-
-
Notifications
You must be signed in to change notification settings - Fork 6.9k
[Feature] Disable Custom Virtual Key Values via UI Setting #23812
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
53d96c8
72aa5fc
c687e63
0b0fe7e
471e0f1
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
|
|
@@ -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: | ||
| 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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Move the 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 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 For Wait, but that means every call to Let me remove this comment and only keep the stricter logic issue (is True vs truthy). Actually, let me reconsider the Actually, looking at the code in ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw)JSON parsing will give proper Python booleans for JSON booleans. So I'll keep this as a style/logic note. Now let me think about other 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.
Let me focus on the most meaningful issues:
Let me also check if there's any issue with the UI change - does the form still submit the Actually, let me check what Actually I think the
Let me look at the Let me also think about race conditions - Let me finalize the review now. |
||
|
|
||
|
|
||
| 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) | ||
|
Comment on lines
+695
to
+696
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
The Moving |
||
|
|
||
| # 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} | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
The 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
Suggested change
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| @pytest.mark.asyncio | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| async def test_generate_service_account_requires_team_id(): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| with pytest.raises(HTTPException): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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"] : []), | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When The fix is to explicitly clear the field when the flag becomes active: useEffect(() => {
if (disableCustomApiKeys) {
form.setFieldValue("key", undefined);
}
}, [disableCustomApiKeys, form]);Add this |
||
| ]} | ||
| /> | ||
| </AccordionBody> | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
is Trueequality checkui_settings.get("disable_custom_api_keys", False) is Truewill only match the Python booleanTrue. If the value is stored/deserialized as the string"true"or the integer1(both of which are truthy), the guard silently passes and the setting is ignored. The standard Python idiom is a simple truthiness check: