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 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
2026-03-07 15:56:43 -08:00
co-authored by Claude Opus 4.6
parent c631708df6
commit af61132a3f
3 changed files with 140 additions and 106 deletions
@@ -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:
@@ -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",
@@ -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())