scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix: enforce disable_custom_api_keys from general_settings (#42437)

Browse source 
* fix: enforce disable_custom_api_keys from general_settings

The gate in _check_custom_key_allowed read the persisted UI settings row
through get_ui_settings_cached, which had two consequences.

A config-file general_settings.disable_custom_api_keys was never enforced,
because the gate only ever looked at the stored ui_settings row. POST
/key/generate with a custom key value returned 200 even with the flag set
to true in config.yaml.

The read went through a DualCache with a 600s TTL, which is per worker
without Redis, and only the worker that served PATCH /update/ui_settings
refreshed it. A gate flipped through the UI was then a coin flip across
workers for up to ten minutes.

Both go away by routing the flag the way the other runtime UI flags are
already routed. Adding it to _RUNTIME_GENERAL_SETTINGS_FLAGS and to the
settings rules' _UI_SETTINGS_FIELDS makes SettingsStore resolve it from the
ui_settings row with the config file winning, and every pod re-reads it on
its own settings sync rather than holding a private cached copy. The two
lists have to stay in step: a flag in one and not the other resolves
against the wrong stored row and silently never reaches a reader, so there
is a test for that invariant.

Writes to the ui_settings table did not publish on the config-sync channel,
so other pods only discovered a change on their next periodic reload. Adding
litellm_uisettings to _CONFIG_SYNCED_TABLE_NAMES puts it on the same pubsub
path model and SSO config writes already use, which cuts cross-pod
propagation from tens of seconds to a few.

The value reaching the gate is run through coerce_bool first. Resolution
hands back the raw YAML value, so a quoted "true" in config.yaml is a str
and the old `is True` check let custom keys straight through.

* test: assert both directions of the coerced config value

* test: assert the runtime flags read back instead of inspecting the registry
This commit is contained in:
yuneng-jiang 2026-09-22 13:28:31 -07:00 • committed by GitHub
parent 2727359a9e
commit 666f6b01b4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 181 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/common_utils/config_sync_pubsub.py
View file

@ -55,6 +55,7 @@ _CONFIG_SYNCED_TABLE_NAMES: Final[frozenset[str]] = frozenset(
"litellm_ssoconfig",
"litellm_cacheconfig",
"litellm_configoverrides",
"litellm_uisettings",
}
)

1
litellm/proxy/config_resolvers/settings_rules.py
View file

@ -48,6 +48,7 @@ _UI_SETTINGS_FIELDS: Final[tuple[str, ...]] = (
"allow_agents_for_team_admins",
"disable_vector_stores_for_internal_users",
"allow_vector_stores_for_team_admins",
"disable_custom_api_keys",
"disable_key_generate_for_org_admin",
"team_admin_editable_team_fields",
)

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

@ -118,9 +118,6 @@ 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.budget_reservation import get_budget_window_start
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,
@ -487,8 +484,10 @@ async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
if custom_key_value is None:
return
ui_settings: Final = await get_ui_settings_cached()
if ui_settings.get("disable_custom_api_keys", False) is True:
from litellm.proxy.config_resolvers.settings_rules import coerce_bool
from litellm.proxy.proxy_server import general_settings
if coerce_bool(general_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,

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

@ -407,6 +407,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
"allow_agents_for_team_admins",
"disable_vector_stores_for_internal_users",
"allow_vector_stores_for_team_admins",
"disable_custom_api_keys",
"disable_key_generate_for_org_admin",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
]

28
tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py
View file

@ -47,6 +47,7 @@ _EXPECTED_CONFIG_SYNCED_TABLE_NAMES = frozenset(
"litellm_proxymodeltable",
"litellm_searchtoolstable",
"litellm_ssoconfig",
"litellm_uisettings",
}
)
@ -702,6 +703,33 @@ async def test_model_repository_write_publishes_via_live_coordination_cache() ->
assert json.loads(message) == {"object_type": "litellm_proxymodeltable"}
async def test_ui_settings_write_publishes_via_live_coordination_cache() -> None:
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import _set_redis_usage_cache
from litellm.repositories.table_repositories import UISettingsRepository
client = _RecordingRedisClient()
prisma_client = MagicMock()
prisma_client.db.litellm_uisettings.upsert = AsyncMock(return_value={"id": "ui_settings"})
table = UISettingsRepository(prisma_client).table
assert isinstance(table, _PublishOnWriteActions)
previous_cache = proxy_server.redis_usage_cache
_set_redis_usage_cache(_FakeRedisCache(client))
try:
await table.upsert(
where={"id": "ui_settings"},
data={"create": {"id": "ui_settings"}, "update": {"ui_settings": "{}"}},
)
finally:
_set_redis_usage_cache(previous_cache)
assert len(client.published) == 1
channel, message = client.published[0]
assert channel == CONFIG_SYNC_CHANNEL
assert json.loads(message) == {"object_type": "litellm_uisettings"}
async def _publish_calls_for_invalidated_param(param_name: str) -> List[Tuple[str, str]]:
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import _set_redis_usage_cache

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

@ -609,10 +609,7 @@ async def test_generate_key_debug_log_never_contains_raw_token(monkeypatch, capl
mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
AsyncMock(return_value={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@ -1494,18 +1491,12 @@ async def test_list_keys_full_object_returns_lifetime_total_spend():
@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={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
# Test with valid new_key
data = RegenerateKeyRequest(new_key="sk-test1234567890abc")
@ -1517,8 +1508,6 @@ async def test_get_new_token_with_valid_key(monkeypatch):
@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
@ -1526,11 +1515,7 @@ async def test_get_new_token_with_invalid_key(monkeypatch):
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={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
# Test with invalid new_key (doesn't start with 'sk-')
data = RegenerateKeyRequest(new_key="invalid-key-123")
@ -1546,8 +1531,6 @@ async def test_get_new_token_with_invalid_key(monkeypatch):
async def test_get_new_token_rejects_short_new_key(monkeypatch):
"""Regression test for LIT-4355: a short custom key like sk-99 must be rejected,
otherwise the stored key_name (sk-...{last 4 chars}) reveals the entire key."""
from unittest.mock import AsyncMock
from fastapi import HTTPException
from litellm.proxy._types import RegenerateKeyRequest
@ -1555,10 +1538,7 @@ async def test_get_new_token_rejects_short_new_key(monkeypatch):
get_new_token,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
AsyncMock(return_value={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
data = RegenerateKeyRequest(new_key="sk-99")
@ -1588,10 +1568,7 @@ async def test_generate_key_fn_rejects_short_custom_key(monkeypatch, short_key):
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
AsyncMock(return_value={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
assert len(short_key) < 16
@ -1628,10 +1605,7 @@ async def test_generate_key_fn_accepts_custom_key_at_minimum_length(monkeypatch)
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
AsyncMock(return_value={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
custom_key = "sk-abcdefghijklm"
assert len(custom_key) == 16
@ -1649,18 +1623,13 @@ async def test_generate_key_fn_accepts_custom_key_at_minimum_length(monkeypatch)
@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}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_custom_api_keys": True})
with pytest.raises(HTTPException) as exc_info:
await _check_custom_key_allowed("sk-custom-key-123")
@ -1672,16 +1641,11 @@ async def test_check_custom_key_allowed_when_disabled(monkeypatch):
@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}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_custom_api_keys": False})
# Should not raise
await _check_custom_key_allowed("sk-custom-key-123")
@ -1690,35 +1654,133 @@ async def test_check_custom_key_allowed_when_enabled(monkeypatch):
@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={}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
# 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
async def test_check_custom_key_allowed_honours_the_config_file(monkeypatch):
"""A config-file general_settings.disable_custom_api_keys is enforced with no stored UI row."""
from fastapi import HTTPException
from litellm.proxy.config_resolvers import SettingsStore
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}),
general_settings = SettingsStore("general_settings")
general_settings.load_yaml({"disable_custom_api_keys": True})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
with pytest.raises(HTTPException) as exc_info:
await _check_custom_key_allowed("sk-custom-key-123456")
assert exc_info.value.status_code == 403
@pytest.mark.parametrize(
("config_value", "blocked"),
[
(True, True),
("true", True),
("True", True),
(1, True),
(False, False),
("false", False),
("False", False),
(0, False),
],
)
@pytest.mark.asyncio
async def test_check_custom_key_allowed_coerces_a_non_bool_config_value(monkeypatch, config_value, blocked):
"""A YAML value that is not a bare bool, such as a quoted "true", still decides the gate."""
from fastapi import HTTPException
from litellm.proxy.config_resolvers import SettingsStore
from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_custom_key_allowed,
)
general_settings = SettingsStore("general_settings")
general_settings.load_yaml({"disable_custom_api_keys": config_value})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
rejected = False
try:
await _check_custom_key_allowed("sk-custom-key-123456")
except HTTPException as e:
rejected = e.status_code == 403
assert rejected is blocked
@pytest.mark.asyncio
async def test_check_custom_key_allowed_config_file_beats_the_stored_ui_row(monkeypatch):
"""The config file owns the flag, so a stored UI row saying false cannot re-open custom keys."""
from fastapi import HTTPException
from litellm.proxy.config_resolvers import SettingsStore
from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_custom_key_allowed,
)
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
apply_runtime_general_settings_flags,
)
general_settings = SettingsStore("general_settings")
general_settings.load_yaml({"disable_custom_api_keys": True})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
apply_runtime_general_settings_flags({"disable_custom_api_keys": False})
with pytest.raises(HTTPException) as exc_info:
await _check_custom_key_allowed("sk-custom-key-123456")
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_check_custom_key_allowed_picks_up_a_ui_write_without_the_serving_pod(monkeypatch):
"""A pod that never served the PATCH enforces the new value after its own settings sync."""
from fastapi import HTTPException
from litellm.proxy.config_resolvers import SettingsStore
from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_custom_key_allowed,
)
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
apply_runtime_general_settings_flags,
)
general_settings = SettingsStore("general_settings")
general_settings.load_yaml({})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
await _check_custom_key_allowed("sk-custom-key-123456")
apply_runtime_general_settings_flags({"disable_custom_api_keys": True})
with pytest.raises(HTTPException) as exc_info:
await _check_custom_key_allowed("sk-custom-key-123456")
assert exc_info.value.status_code == 403
@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 litellm.proxy.management_endpoints.key_management_endpoints import (
_check_custom_key_allowed,
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_custom_api_keys": True})
# Should not raise — None means auto-generate
await _check_custom_key_allowed(None)
@ -1726,8 +1788,6 @@ async def test_check_custom_key_allowed_none_key_always_passes(monkeypatch):
@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
@ -1735,10 +1795,7 @@ async def test_get_new_token_rejected_when_custom_keys_disabled(monkeypatch):
get_new_token,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
AsyncMock(return_value={"disable_custom_api_keys": True}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_custom_api_keys": True})
data = RegenerateKeyRequest(new_key="sk-custom-regen-key")
@ -1751,17 +1808,12 @@ async def test_get_new_token_rejected_when_custom_keys_disabled(monkeypatch):
@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}),
)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_custom_api_keys": True})
data = RegenerateKeyRequest() # no new_key
result = await get_new_token(data)

26
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
View file

@ -3998,6 +3998,32 @@ class TestSyncUiSettingsToGeneralSettings:
assert general_settings["forward_client_headers_to_llm_api"] is True
assert general_settings.source("forward_client_headers_to_llm_api") == "db"
def test_every_runtime_flag_reaches_a_reader_once_applied(self, monkeypatch):
"""A flag the settings rules do not route to the ui_settings row is stored but never read back."""
from litellm.proxy import proxy_server
from litellm.proxy.config_resolvers import SettingsStore
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
_RUNTIME_GENERAL_SETTINGS_FLAGS,
apply_runtime_general_settings_flags,
)
general_settings = SettingsStore("general_settings")
general_settings.load_yaml({})
monkeypatch.setattr(proxy_server, "general_settings", general_settings)
stored = {
key: (["tpm_limit"] if key == TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING else True)
for key in _RUNTIME_GENERAL_SETTINGS_FLAGS
}
assert stored
apply_runtime_general_settings_flags(stored)
read_back = {key: general_settings.get(key) for key in stored}
assert read_back == stored
def test_applied_runtime_flags_cannot_override_the_config_file(self, monkeypatch):
from litellm.proxy import proxy_server
from litellm.proxy.config_resolvers import SettingsStore