From af61132a3f4331a2ff352d38ec92b03835a126e8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 7 Mar 2026 15:56:43 -0800 Subject: [PATCH] refactor: use DualCache for UI settings reads and get_team_object for team lookup - Add get_ui_settings_cached() helper that reads from DualCache first, falls back to DB, and populates cache on miss. - Update update_ui_settings() to set cache after DB write so subsequent reads see new values immediately. - Replace raw prisma_client.db.litellm_teamtable.find_unique with the existing get_team_object helper which uses DualCache. - Update all tests to mock get_ui_settings_cached and get_team_object instead of raw DB calls. Co-Authored-By: Claude Opus 4.6 --- .../internal_user_endpoints.py | 54 ++++--- .../proxy_setting_endpoints.py | 52 +++++++ .../test_internal_user_endpoints.py | 140 ++++++++---------- 3 files changed, 140 insertions(+), 106 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d43efd8138..8d9da86aab 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -30,7 +30,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, get_daily_activity_aggregated, ) -from litellm.proxy.auth.auth_checks import get_user_object +from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -1869,21 +1869,17 @@ async def ui_view_users( proxy_logging_obj, user_api_key_cache, ) + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + get_ui_settings_cached, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) try: - # Read the scope_user_search_to_org flag from the DB - ui_settings_row = ( - await prisma_client.db.litellm_uisettings.find_unique( - where={"id": "ui_settings"} - ) - ) - scope_flag = False - if ui_settings_row is not None: - settings_json = ui_settings_row.settings or {} # type: ignore[union-attr] - scope_flag = bool(settings_json.get("scope_user_search_to_org", False)) + # Read the scope_user_search_to_org flag (cached) + ui_settings = await get_ui_settings_cached() + scope_flag = bool(ui_settings.get("scope_user_search_to_org", False)) org_filter_ids: Optional[List[str]] = None @@ -1915,27 +1911,29 @@ async def ui_view_users( if org_admin_org_ids: org_filter_ids = org_admin_org_ids elif team_id is not None: - # Look up the team to check if it belongs to an org - team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - if team_row is not None: - team_obj = LiteLLM_TeamTable(**team_row.model_dump()) - if _is_user_team_admin(user_api_key_dict, team_obj): - if team_obj.organization_id: - org_filter_ids = [team_obj.organization_id] - else: - raise HTTPException( - status_code=403, - detail={ - "error": "scope_user_search_to_org is enabled and this team is not part of an organization. Contact your proxy admin to adjust this setting." - }, - ) + # Look up the team via cached helper + try: + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + raise HTTPException( + status_code=403, + detail={ + "error": "scope_user_search_to_org is enabled. Only proxy admins, organization admins, or team admins can search users." + }, + ) + if _is_user_team_admin(user_api_key_dict, team_obj): + if team_obj.organization_id: + org_filter_ids = [team_obj.organization_id] else: raise HTTPException( status_code=403, detail={ - "error": "scope_user_search_to_org is enabled. Only proxy admins, organization admins, or team admins can search users." + "error": "scope_user_search_to_org is enabled and this team is not part of an organization. Contact your proxy admin to adjust this setting." }, ) else: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 9603f55e48..04b524ebbb 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -980,6 +980,48 @@ async def get_in_product_nudges(): return InProductNudgeResponse(is_claude_code_enabled=False) +UI_SETTINGS_CACHE_KEY = "ui_settings:settings_dict" + + +async def get_ui_settings_cached() -> Dict[str, Any]: + """ + Return the persisted UI settings dict, using DualCache for reads. + + Cache hit → return cached dict immediately. + Cache miss → read from DB, populate cache, return dict. + """ + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + # 1. Try cache + cached = await user_api_key_cache.async_get_cache(key=UI_SETTINGS_CACHE_KEY) + if cached is not None and isinstance(cached, dict): + return cached + + # 2. Fallback to DB + if prisma_client is None: + return {} + + db_record = await prisma_client.db.litellm_uisettings.find_unique( + where={"id": "ui_settings"} + ) + ui_settings: Dict[str, Any] = {} + if db_record and db_record.ui_settings: + raw = db_record.ui_settings + ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) + + # Sanitize + ui_settings = { + k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS + } + + # 3. Populate cache + await user_api_key_cache.async_set_cache( + key=UI_SETTINGS_CACHE_KEY, value=ui_settings + ) + + return ui_settings + + @router.get( "/get/ui_settings", tags=["UI Settings"], @@ -1108,6 +1150,16 @@ async def update_ui_settings( general_settings.update(_flags_to_sync) + # Invalidate + set DualCache so subsequent reads see the new values immediately + from litellm.proxy.proxy_server import user_api_key_cache + + sanitized = { + k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS + } + await user_api_key_cache.async_set_cache( + key=UI_SETTINGS_CACHE_KEY, value=sanitized + ) + return { "message": "UI settings updated successfully", "status": "success", diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 6316d13b98..16d16c040a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -54,11 +54,11 @@ async def test_ui_view_users_with_null_email(mocker, caplog): mock_prisma_client.db.litellm_usertable.find_many = mock_find_many - # Flag OFF by default — no settings row - async def mock_find_unique_settings(*args, **kwargs): - return None - - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + # Flag OFF by default + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={}, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -92,10 +92,10 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): mock_prisma_client.db.litellm_usertable.find_many = mock_find_many # Flag OFF by default - async def mock_find_unique_settings(*args, **kwargs): - return None - - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={}, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) await ui_view_users( @@ -132,13 +132,10 @@ async def test_ui_view_users_org_admin_filtered_by_org(mocker): mock_prisma_client.db.litellm_usertable.find_many = mock_find_many # Flag ON - mock_settings_row = mocker.MagicMock() - mock_settings_row.settings = {"scope_user_search_to_org": True} - - async def mock_find_unique_settings(*args, **kwargs): - return mock_settings_row - - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock()) @@ -185,13 +182,10 @@ async def test_ui_view_users_non_org_admin_returns_403(mocker): mock_prisma_client = mocker.MagicMock() # Flag ON - mock_settings_row = mocker.MagicMock() - mock_settings_row.settings = {"scope_user_search_to_org": True} - - async def mock_find_unique_settings(*args, **kwargs): - return mock_settings_row - - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock()) @@ -237,11 +231,11 @@ async def test_ui_view_users_flag_off_internal_user_can_search(mocker): mock_prisma_client.db.litellm_usertable.find_many = mock_find_many - # Flag OFF — no settings row - async def mock_find_unique_settings(*args, **kwargs): - return None - - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + # Flag OFF + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={}, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) response = await ui_view_users( @@ -261,7 +255,7 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker): """ Flag ON, team admin for org-bound team: org filter is applied using team's org. """ - from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy._types import LiteLLM_TeamTableCachedObj mock_prisma_client = mocker.MagicMock() org_id = "org-456" @@ -278,30 +272,26 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker): mock_prisma_client.db.litellm_usertable.find_many = mock_find_many # Flag ON - mock_settings_row = mocker.MagicMock() - mock_settings_row.settings = {"scope_user_search_to_org": True} + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) - async def mock_find_unique_settings(*args, **kwargs): - return mock_settings_row + # Mock get_team_object + team_obj = LiteLLM_TeamTableCachedObj( + team_id=tid, + team_alias="test-team", + organization_id=org_id, + members_with_roles=[{"user_id": "team-admin-user", "role": "admin"}], + ) - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + async def mock_get_team_object(*args, **kwargs): + return team_obj - # Team lookup - mock_team_row = mocker.MagicMock() - mock_team_row.model_dump.return_value = { - "team_id": tid, - "team_alias": "test-team", - "organization_id": org_id, - "members_with_roles": [{"user_id": "team-admin-user", "role": "admin"}], - "admins": [], - "members": [], - "blocked": False, - } - - async def mock_find_unique_team(*args, **kwargs): - return mock_team_row - - mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_team_object", + side_effect=mock_get_team_object, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock()) @@ -337,35 +327,32 @@ async def test_ui_view_users_flag_on_team_admin_non_org_team_403(mocker): Flag ON, team admin for non-org team: returns 403. """ from fastapi import HTTPException + from litellm.proxy._types import LiteLLM_TeamTableCachedObj mock_prisma_client = mocker.MagicMock() tid = "team-no-org" # Flag ON - mock_settings_row = mocker.MagicMock() - mock_settings_row.settings = {"scope_user_search_to_org": True} + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) - async def mock_find_unique_settings(*args, **kwargs): - return mock_settings_row + # Mock get_team_object — team has no organization_id + team_obj = LiteLLM_TeamTableCachedObj( + team_id=tid, + team_alias="no-org-team", + organization_id=None, + members_with_roles=[{"user_id": "team-admin-user", "role": "admin"}], + ) - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + async def mock_get_team_object(*args, **kwargs): + return team_obj - # Team lookup — no organization_id - mock_team_row = mocker.MagicMock() - mock_team_row.model_dump.return_value = { - "team_id": tid, - "team_alias": "no-org-team", - "organization_id": None, - "members_with_roles": [{"user_id": "team-admin-user", "role": "admin"}], - "admins": [], - "members": [], - "blocked": False, - } - - async def mock_find_unique_team(*args, **kwargs): - return mock_team_row - - mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_team_object", + side_effect=mock_get_team_object, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock()) @@ -409,13 +396,10 @@ async def test_ui_view_users_flag_on_non_admin_no_team_id_403(mocker): mock_prisma_client = mocker.MagicMock() # Flag ON - mock_settings_row = mocker.MagicMock() - mock_settings_row.settings = {"scope_user_search_to_org": True} - - async def mock_find_unique_settings(*args, **kwargs): - return mock_settings_row - - mock_prisma_client.db.litellm_uisettings.find_unique = mock_find_unique_settings + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())