fix(proxy): scope dashboard search grants to user (#48464)

Co-authored-by: Jodok Batlogg <jodok@batlogg.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/auth/auth_checks.py

@@ -6461,8 +6461,29 @@ def _search_tool_deny_by_default(general_settings: Mapping[str, object]) -> bool
         return True
 
 
+def _is_dashboard_session(valid_token: UserAPIKeyAuth | None) -> bool:
+    return valid_token is not None and valid_token.team_id == UI_TEAM_ID
+
+
+def _is_search_tool_strict_identity(valid_token: UserAPIKeyAuth) -> bool:
+    if not _is_dashboard_session(valid_token):
+        return _is_strict_grant_identity(valid_token)
+    return valid_token.user_role not in (
+        LitellmUserRoles.PROXY_ADMIN,
+        LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
+    )
+
+
+def _search_tool_strict_layers(
+    valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None
+) -> frozenset[GrantLayer]:
+    if _is_dashboard_session(valid_token):
+        return frozenset(("user",))
+    return _strict_grant_layers(True, valid_token, team_object)
+
+
 def is_search_tool_deny_by_default_applied(valid_token: UserAPIKeyAuth, general_settings: Mapping[str, object]) -> bool:
-    return _search_tool_deny_by_default(general_settings) and _is_strict_grant_identity(valid_token)
+    return _search_tool_deny_by_default(general_settings) and _is_search_tool_strict_identity(valid_token)
 
 
 def _search_tool_denied(object_type: GrantLayer, message: str) -> ProxyException:
@@ -6506,12 +6527,12 @@ async def resolve_search_tool_grants(
     load_team_object: TeamObjectLoader,
 ) -> SearchToolGrants:
     """
-    Without `general_settings.search_tool_deny_by_default`, or for the master key and dashboard sessions,
+    Without `general_settings.search_tool_deny_by_default`, or for the master key and admin dashboard sessions,
     the key and team lists only restrict when nonempty. With it, the identities `_strict_grant_layers` names
     must each list the tool, and a missing record, `null`, `[]`, or a team or user that fails to load grants nothing
     """
     if not is_search_tool_deny_by_default_applied(valid_token, general_settings):
-        team_object: Final = await load_team_object()
+        team_object: Final = await load_team_object() if not _is_dashboard_session(valid_token) else None
         return SearchToolGrants(
             grants=MappingProxyType(
                 {"key": valid_token.object_permission, "team": team_object.object_permission if team_object else None}
@@ -6521,8 +6542,10 @@ async def resolve_search_tool_grants(
 
     from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
 
-    strict_team_object: Final = await _strict_team_object(load_team_object)
-    strict_layers: Final = _strict_grant_layers(True, valid_token, strict_team_object)
+    strict_team_object: Final = (
+        await _strict_team_object(load_team_object) if not _is_dashboard_session(valid_token) else None
+    )
+    strict_layers: Final = _search_tool_strict_layers(valid_token, strict_team_object)
     if prisma_client is None:
         return SearchToolGrants(grants=MappingProxyType({}), strict_layers=strict_layers)
     user_object: Final = (
@@ -6607,7 +6630,7 @@ def check_unregistered_search_fallback(
     """The unregistered provider fallback names no search tool that a grant could list, so deny-by-default denies it"""
     if not is_search_tool_deny_by_default_applied(valid_token, general_settings):
         return True
-    layer: Final = min(_strict_grant_layers(True, valid_token, None), key=("key", "team", "user").index)
+    layer: Final = min(_search_tool_strict_layers(valid_token, None), key=("key", "team", "user").index)
     raise _search_tool_denied(
         layer,
         "No registered search tool is available and search_tool_deny_by_default is enabled",

tests/unit/proxy/auth/test_search_tool_deny_by_default.py

@@ -1,8 +1,9 @@
 from collections.abc import Mapping
-from typing import Final, TypedDict
+from typing import Final, TypedDict, cast
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from fastapi import HTTPException
 
 from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
 from litellm.proxy import proxy_server
@@ -78,6 +79,14 @@ def _team_loader(team_tools: list[str] | None | str) -> TeamObjectLoader:
     return load
 
 
+def _missing_dashboard_team_loader() -> tuple[TeamObjectLoader, AsyncMock]:
+    missing_team: Final = HTTPException(
+        status_code=404, detail="Team doesn't exist in db. Team=litellm-dashboard"
+    )
+    load_team: Final = AsyncMock(side_effect=missing_team)
+    return cast(TeamObjectLoader, load_team), load_team
+
+
 async def _no_team() -> None:
     return None
 
@@ -194,9 +203,6 @@ async def test_search_tool_access_is_unchanged_without_search_tool_deny_by_defau
             id="proxy admin virtual key not exempt",
         ),
         pytest.param(CallerFields(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, None, None, id="master key exempt"),
-        pytest.param(
-            CallerFields(team_id=UI_TEAM_ID, virtual_key=False), None, None, None, id="dashboard session exempt"
-        ),
     ],
 )
 @pytest.mark.asyncio
@@ -248,6 +254,79 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off(
     get_user_object.assert_not_awaited()
 
 
+@pytest.mark.parametrize(
+    "general_settings",
+    [{}, {"search_tool_deny_by_default": False}],
+    ids=["omitted", "false"],
+)
+@pytest.mark.parametrize(
+    "user_role",
+    [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER],
+    ids=["admin", "internal-user"],
+)
+@pytest.mark.parametrize(
+    "key_tools, expected",
+    [
+        pytest.param(None, None, id="unrestricted-key"),
+        pytest.param(["search-a"], None, id="key-grants-search-a"),
+        pytest.param(["search-b"], _LEGACY, id="key-excludes-search-a"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_dashboard_search_tool_grants_without_deny_by_default(
+    general_settings: Mapping[str, object],
+    user_role: LitellmUserRoles,
+    key_tools: list[str] | None,
+    expected: ProxyErrorTypes | None,
+) -> None:
+    caller: Final = _caller(
+        virtual_key=False,
+        key_tools=key_tools,
+        team_id=UI_TEAM_ID,
+        user_role=user_role,
+    )
+    load_team, load_team_mock = _missing_dashboard_team_loader()
+    assert await _denied_by(general_settings, caller, load_team) == expected
+    load_team_mock.assert_not_awaited()
+
+
+@pytest.mark.parametrize(
+    "user_role",
+    [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY],
+    ids=["admin", "admin-view-only"],
+)
+@pytest.mark.asyncio
+async def test_admin_dashboard_search_tool_grants_with_deny_by_default(user_role: LitellmUserRoles) -> None:
+    caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=user_role)
+    load_team, load_team_mock = _missing_dashboard_team_loader()
+    assert await _denied_by(_DENY_ON, caller, load_team) is None
+    load_team_mock.assert_not_awaited()
+
+
+@pytest.mark.parametrize(
+    "user_id, user_tools, expected",
+    [
+        pytest.param("user-1", ["search-a"], None, id="user-grants-search-a"),
+        pytest.param("user-1", ["search-b"], _USER, id="user-grants-another-tool"),
+        pytest.param("user-1", None, _USER, id="user-grant-null"),
+        pytest.param("user-2", None, _USER, id="no-user-row"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_internal_dashboard_search_tool_requires_the_user_grant_with_deny_by_default(
+    cache: UserApiKeyCache,
+    user_id: str,
+    user_tools: list[str] | None,
+    expected: ProxyErrorTypes | None,
+) -> None:
+    if user_id == "user-1":
+        _cache_user(cache, user_tools)
+    caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_id=user_id)
+    load_team, load_team_mock = _missing_dashboard_team_loader()
+    assert await _denied_by(_DENY_ON, caller, load_team) == expected
+    load_team_mock.assert_not_awaited()
+
+
 @pytest.mark.parametrize(
     "general_settings, caller, expected",
     [
@@ -257,6 +336,24 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off(
         pytest.param(_DENY_ON, _caller(virtual_key=False), _USER, id="keyless user"),
         pytest.param(_DENY_ON, _caller(user_role=LitellmUserRoles.PROXY_ADMIN), _KEY, id="proxy admin key"),
         pytest.param(_DENY_ON, _caller(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, id="master key"),
+        pytest.param(
+            _DENY_ON,
+            _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.INTERNAL_USER),
+            _USER,
+            id="internal dashboard user",
+        ),
+        pytest.param(
+            _DENY_ON,
+            _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN),
+            None,
+            id="admin dashboard",
+        ),
+        pytest.param(
+            _DENY_ON,
+            _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
+            None,
+            id="admin view-only dashboard",
+        ),
     ],
 )
 def test_unregistered_search_fallback_follows_search_tool_deny_by_default(
@@ -268,3 +365,12 @@ def test_unregistered_search_fallback_follows_search_tool_deny_by_default(
     with pytest.raises(ProxyException) as exc_info:
         check_unregistered_search_fallback(caller, general_settings)
     assert (exc_info.value.type, exc_info.value.code) == (expected, "403")
+
+
+@pytest.mark.asyncio
+async def test_search_preserves_missing_real_team_error_without_deny_by_default() -> None:
+    missing_team: Final = HTTPException(status_code=404, detail="Team doesn't exist in db. Team=team-1")
+    load_team: Final = AsyncMock(side_effect=missing_team)
+    with pytest.raises(HTTPException) as exc_info:
+        await can_caller_call_search_tool("search-a", _caller(team_id="team-1"), {}, load_team)
+    assert exc_info.value is missing_team