From 666f6b01b4c4dc474bbb4d9d129c127cfada6961 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 22 Sep 2026 13:28:31 -0700 Subject: [PATCH] fix: enforce disable_custom_api_keys from general_settings (#42437) * 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 --- .../proxy/common_utils/config_sync_pubsub.py | 1 + .../proxy/config_resolvers/settings_rules.py | 1 + .../key_management_endpoints.py | 9 +- .../proxy_setting_endpoints.py | 1 + .../common_utils/test_config_sync_pubsub.py | 28 +++ .../test_key_management_endpoints.py | 188 +++++++++++------- .../test_proxy_setting_endpoints.py | 26 +++ 7 files changed, 181 insertions(+), 73 deletions(-) diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index 6d781babe63..b20c0d9c9a5 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -55,6 +55,7 @@ _CONFIG_SYNCED_TABLE_NAMES: Final[frozenset[str]] = frozenset( "litellm_ssoconfig", "litellm_cacheconfig", "litellm_configoverrides", + "litellm_uisettings", } ) diff --git a/litellm/proxy/config_resolvers/settings_rules.py b/litellm/proxy/config_resolvers/settings_rules.py index f346dd6198d..1f0adfc5248 100644 --- a/litellm/proxy/config_resolvers/settings_rules.py +++ b/litellm/proxy/config_resolvers/settings_rules.py @@ -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", ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fbc7cf18003..801c45501ce 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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, diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index f37b82f3b54..d520177965c 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -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, ] diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 2a0ebf9492a..0f64ef2b4ca 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8eaa4901c59..c539a10e1e2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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) diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index f40848bc38b..0140fcaba21 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -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