From cacf3debac4f3043f81f6458c445ff054869f418 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 12:07:09 -0700 Subject: [PATCH] fix(proxy): scope a dashboard session on /search_tools/list to its user's real teams Treating the reserved litellm-dashboard team id as "no team" dropped team-level scoping for every Admin UI session, so the Search Tools page listed every tool on the proxy, including ids and api_base for tools the caller cannot invoke. The listing now resolves the sentinel to the real teams backing the session user and shows a tool when at least one of them allowlists it, matching POST /search. Both surfaces share one resolver and one predicate in search_tool_access, which retires the duplicate team scoping in auth_checks and search_tool_management. --- litellm/proxy/auth/auth_checks.py | 23 --- .../search_tool_management.py | 82 ++++------ .../test_search_tool_management.py | 153 +++++++++++++++--- 3 files changed, 156 insertions(+), 102 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3fba464dd23..c24b891d812 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3538,29 +3538,6 @@ async def can_team_call_search_tool( ) -async def can_user_view_search_tool( - search_tool_name: str, - valid_token: UserAPIKeyAuth, - team_object: LiteLLM_TeamTable | None, -) -> bool: - """ - Boolean variant of the key + team authorization enforced on /search, used to - scope /search_tools/list so a non-admin caller only sees tools it may invoke. - """ - try: - await can_key_call_search_tool( - search_tool_name=search_tool_name, - valid_token=valid_token, - ) - await can_team_call_search_tool( - search_tool_name=search_tool_name, - team_object=team_object, - ) - except ProxyException: - return False - return True - - async def is_valid_fallback_model( model: str, llm_router: Router | None, diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 69edf681e4d..f19ae2721ad 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -2,21 +2,29 @@ CRUD ENDPOINTS FOR SEARCH TOOLS """ -from collections.abc import Awaitable, Callable +import asyncio +from collections.abc import Sequence from datetime import datetime -from typing import Any, Final, TypeAlias +from typing import Any, Final from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from litellm._logging import verbose_proxy_logger -from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import ( LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.search_endpoints.search_tool_access import ( + SessionTeamIdsLookup, + TeamObjectLookup, + can_view_search_tool, + resolve_allowlist_teams, + session_team_ids_from_db, + team_object_from_db, +) from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry from litellm.types.search import ( ListSearchToolsResponse, @@ -48,46 +56,26 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None: return value -TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]] - - -async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: - from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, +async def _is_visible( + tool: SearchToolInfoResponse, + user_api_key_dict: UserAPIKeyAuth, + teams: Sequence[LiteLLM_TeamTable], +) -> bool: + tool_name: Final = tool.get("search_tool_name") + if not tool_name: + return False + return await can_view_search_tool( + search_tool_name=tool_name, + user_api_key_dict=user_api_key_dict, + teams=teams, ) - return await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - -def _allowlist_team_id(user_api_key_dict: UserAPIKeyAuth) -> str | None: - """ - The team whose object_permission allowlist scopes this caller, or None when there is none. - - Every Admin UI session key is stamped with UI_SESSION_TOKEN_TEAM_ID, a reserved sentinel that - never has a row in LiteLLM_TeamTable (`/team/new` rejects it as a real team id), so looking it - up would raise 404 instead of resolving a team. It carries no allowlist of its own, so the - caller is scoped by its key-level allowlist alone. Any other team id is looked up for real and - a failed lookup still surfaces. - """ - team_id: Final = user_api_key_dict.team_id - if not team_id or team_id == UI_SESSION_TOKEN_TEAM_ID: - return None - return team_id - async def _filter_visible_search_tools( search_tools: list[SearchToolInfoResponse], user_api_key_dict: UserAPIKeyAuth, - lookup_team_object: TeamObjectLookup = _team_object_from_db, + lookup_team_object: TeamObjectLookup = team_object_from_db, + lookup_session_team_ids: SessionTeamIdsLookup = session_team_ids_from_db, ) -> list[SearchToolInfoResponse]: """ Drop search tools the caller is not authorized to invoke, applying the same @@ -99,23 +87,9 @@ async def _filter_visible_search_tools( ): return search_tools - from litellm.proxy.auth.auth_checks import can_user_view_search_tool - - allowlist_team_id: Final = _allowlist_team_id(user_api_key_dict) - team_object: Final[LiteLLM_TeamTable | None] = ( - await lookup_team_object(allowlist_team_id, user_api_key_dict) if allowlist_team_id else None - ) - - visible: Final[list[SearchToolInfoResponse]] = [] - for tool in search_tools: - tool_name = tool.get("search_tool_name") - if tool_name and await can_user_view_search_tool( - search_tool_name=tool_name, - valid_token=user_api_key_dict, - team_object=team_object, - ): - visible.append(tool) - return visible + teams: Final = await resolve_allowlist_teams(user_api_key_dict, lookup_team_object, lookup_session_team_ids) + visible: Final = await asyncio.gather(*(_is_visible(tool, user_api_key_dict, teams) for tool in search_tools)) + return [tool for tool, is_visible in zip(search_tools, visible) if is_visible] @router.get( diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index a64397d9818..f20647db4e1 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -15,6 +15,7 @@ sys.path.insert( from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth, ) @@ -767,6 +768,20 @@ async def test_list_search_tools_unrestricted_internal_user_sees_all(): assert names == {"db-tool-1", "db-tool-2", "db-tool-3"} +def _session_user_on(*team_ids: str) -> AsyncMock: + return AsyncMock(return_value=LiteLLM_UserTable(user_id="internal_user", teams=list(team_ids))) + + +def _team_allowing(team_id: str, search_tools: list[str]) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"op-{team_id}", + search_tools=search_tools, + ), + ) + + @pytest.mark.asyncio async def test_list_search_tools_scoped_to_team_object_permission(): """A team-level search_tools allowlist also scopes the listing for a non-admin caller.""" @@ -775,19 +790,12 @@ async def test_list_search_tools_scoped_to_team_object_permission(): user_id="internal_user", team_id="team-1", ) - team_object = LiteLLM_TeamTable( - team_id="team-1", - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="op-team", - search_tools=["db-tool-2"], - ), - ) with ( _mock_search_tool_backend(_scoping_db_tools()), patch( "litellm.proxy.auth.auth_checks.get_team_object", - AsyncMock(return_value=team_object), + AsyncMock(return_value=_team_allowing("team-1", ["db-tool-2"])), ), _override_auth(team_member), ): @@ -839,12 +847,47 @@ def _team_ids_looked_up(lookup: AsyncMock) -> list[str]: @pytest.mark.asyncio -async def test_list_search_tools_dashboard_session_key_does_not_look_up_the_ui_team(): +async def test_list_search_tools_dashboard_session_scoped_to_the_users_real_team_allowlist(): """ Regression: the Admin UI session key is stamped with the reserved team id - ``litellm-dashboard``, which has no row in LiteLLM_TeamTable. Resolving it as a real team - raised 404, which the endpoint reported as a 500, so the Search Tools page was broken for - every non-admin browsing the dashboard. + ``litellm-dashboard``, which has no row in LiteLLM_TeamTable. Treating that as "no team" + dropped team-level scoping, so the Search Tools page listed every tool on the proxy, + including ids and api_base for tools the caller cannot invoke. + """ + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + dashboard_session_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=UI_SESSION_TOKEN_TEAM_ID, + ) + + with ( + _mock_search_tool_backend(_scoping_db_tools()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + AsyncMock(return_value=_team_allowing("team-1", ["db-tool-2"])), + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + _session_user_on("team-1"), + ), + _override_auth(dashboard_session_user), + ): + response = TestClient(app).get("/search_tools/list") + + assert response.status_code == 200 + tools = response.json()["search_tools"] + assert [t["search_tool_name"] for t in tools] == ["db-tool-2"] + leaked = {t["litellm_params"].get("api_base") for t in tools} + assert "https://api.perplexity.ai" not in leaked + + +@pytest.mark.asyncio +async def test_list_search_tools_dashboard_session_never_resolves_the_ui_team(): + """ + Regression: resolving ``litellm-dashboard`` as a real team raised 404, which the endpoint + reported as a 500, so the Search Tools page was broken for every non-admin. """ from litellm.constants import UI_SESSION_TOKEN_TEAM_ID @@ -866,6 +909,10 @@ async def test_list_search_tools_dashboard_session_key_does_not_look_up_the_ui_t "litellm.proxy.auth.auth_checks.get_team_object", ui_team_is_not_a_real_team, ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + _session_user_on(), + ), _override_auth(dashboard_session_user), ): response = TestClient(app).get("/search_tools/list") @@ -879,8 +926,8 @@ async def test_list_search_tools_dashboard_session_key_does_not_look_up_the_ui_t @pytest.mark.asyncio async def test_filter_visible_search_tools_dashboard_session_still_honors_key_allowlist(): """ - Skipping the synthetic team must not widen visibility: a dashboard session whose key - carries a search_tools allowlist stays scoped to it. + Resolving the session through its user's teams must not widen visibility: a dashboard + session whose key carries a search_tools allowlist stays scoped to it. """ from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.search_endpoints.search_tool_management import ( @@ -896,16 +943,47 @@ async def test_filter_visible_search_tools_dashboard_session_still_honors_key_al search_tools=["db-tool-3"], ), ) - lookup = AsyncMock() + lookup = AsyncMock(return_value=_team_allowing("team-1", [])) visible = await _filter_visible_search_tools( _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), restricted_dashboard_session, lookup, + AsyncMock(return_value=["team-1"]), ) assert [t["search_tool_name"] for t in visible] == ["db-tool-3"] - lookup.assert_not_awaited() + assert _team_ids_looked_up(lookup) == ["team-1"] + + +@pytest.mark.asyncio +async def test_filter_visible_search_tools_dashboard_session_sees_the_union_of_its_teams(): + """A user on several teams sees every tool any one of those teams allowlists.""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + dashboard_session_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=UI_SESSION_TOKEN_TEAM_ID, + ) + lookup = AsyncMock( + side_effect=[ + _team_allowing("team-1", ["db-tool-1"]), + _team_allowing("team-2", ["db-tool-3"]), + ] + ) + + visible = await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), + dashboard_session_user, + lookup, + AsyncMock(return_value=["team-1", "team-2"]), + ) + + assert [t["search_tool_name"] for t in visible] == ["db-tool-1", "db-tool-3"] @pytest.mark.asyncio @@ -920,15 +998,7 @@ async def test_filter_visible_search_tools_still_applies_a_real_team_allowlist() user_id="internal_user", team_id="team-1", ) - lookup = AsyncMock( - return_value=LiteLLM_TeamTable( - team_id="team-1", - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="op-team", - search_tools=["db-tool-2"], - ), - ) - ) + lookup = AsyncMock(return_value=_team_allowing("team-1", ["db-tool-2"])) visible = await _filter_visible_search_tools( _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), @@ -997,3 +1067,36 @@ async def test_list_search_tools_reports_a_missing_real_team_as_404(): assert response.status_code == 404 assert "search_tools" not in response.json() + + +@pytest.mark.asyncio +async def test_list_search_tools_dashboard_session_shows_nothing_when_no_team_loads(): + """ + Regression: a dashboard session whose every team failed to load resolved to an empty team set, + which reads as "belongs to no team" and skipped team scoping, so the listing showed every tool + on the proxy to a user entitled to none of them. + """ + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + dashboard_session_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=UI_SESSION_TOKEN_TEAM_ID, + ) + every_team_is_gone = AsyncMock( + side_effect=HTTPException( + status_code=404, + detail={"error": "Team doesn't exist in db. Team=team-1."}, + ) + ) + + with ( + _mock_search_tool_backend(_scoping_db_tools()), + patch("litellm.proxy.auth.auth_checks.get_team_object", every_team_is_gone), + patch("litellm.proxy.auth.auth_checks.get_user_object", _session_user_on("team-1")), + _override_auth(dashboard_session_user), + ): + response = TestClient(app).get("/search_tools/list") + + assert response.status_code == 404 + assert every_team_is_gone.await_count == 1