From 54e9964eb829b242db45d99d3f6214262b7ae56f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 5 Aug 2026 23:22:11 -0700 Subject: [PATCH] fix(proxy): stop resolving the UI session sentinel team on /search_tools/list Every Admin UI session key is stamped with the reserved team id `litellm-dashboard`, which never has a row in LiteLLM_TeamTable, so the team lookup in _filter_visible_search_tools raised 404 and the endpoint returned 500 for every non-admin dashboard session. Skip the lookup for that sentinel and scope the caller by its key-level allowlist alone, matching how MCP and agent permission checks already treat it. A real team id is still resolved, and a genuine lookup failure now surfaces with its own status instead of being masked as a 500. --- .../search_tool_management.py | 66 ++++-- .../test_search_tool_management.py | 194 ++++++++++++++++++ 2 files changed, 241 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 381a5e14ca2..69edf681e4d 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -2,13 +2,15 @@ CRUD ENDPOINTS FOR SEARCH TOOLS """ +from collections.abc import Awaitable, Callable from datetime import datetime -from typing import Any, Final +from typing import Any, Final, TypeAlias 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, @@ -46,9 +48,46 @@ 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, + ) + + 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, ) -> list[SearchToolInfoResponse]: """ Drop search tools the caller is not authorized to invoke, applying the same @@ -60,25 +99,12 @@ async def _filter_visible_search_tools( ): return search_tools - from litellm.proxy.auth.auth_checks import ( - can_user_view_search_tool, - get_team_object, - ) - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.auth.auth_checks import can_user_view_search_tool - team_object: LiteLLM_TeamTable | None = None - if user_api_key_dict.team_id: - team_object = await get_team_object( - team_id=user_api_key_dict.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, - ) + 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: @@ -213,6 +239,8 @@ async def list_search_tools( visible_search_tools: Final = await _filter_visible_search_tools(search_tool_configs, user_api_key_dict) return ListSearchToolsResponse(search_tools=visible_search_tools) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception("Error getting search tools: %s", e) raise HTTPException(status_code=500, detail=str(e)) 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 f2ccfcd0155..158dee7cb78 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 @@ -5,6 +5,7 @@ from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient sys.path.insert( @@ -815,3 +816,196 @@ async def test_list_search_tools_admin_with_restricted_key_still_sees_all(): assert response.status_code == 200 names = {t["search_tool_name"] for t in response.json()["search_tools"]} assert names == {"db-tool-1", "db-tool-2", "db-tool-3"} + + +def _search_tool_responses(*names): + from litellm.types.search import SearchToolInfoResponse + + return [ + SearchToolInfoResponse( + search_tool_id=f"id-{name}", + search_tool_name=name, + litellm_params={"search_provider": "perplexity"}, + search_tool_info=None, + created_at=None, + updated_at=None, + is_from_config=False, + ) + for name in names + ] + + +def _recording_team_lookup(team_object=None, raises=None): + """A `TeamObjectLookup` double that records the team ids it was asked to resolve.""" + asked_for = [] + + async def _lookup(team_id, user_api_key_dict): + asked_for.append(team_id) + if raises is not None: + raise raises + return team_object + + return _lookup, asked_for + + +@pytest.mark.asyncio +async def test_list_search_tools_dashboard_session_key_does_not_look_up_the_ui_team(): + """ + 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. + """ + 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, + ) + ui_team_is_not_a_real_team = AsyncMock( + side_effect=HTTPException( + status_code=404, + detail={"error": f"Team doesn't exist in db. Team={UI_SESSION_TOKEN_TEAM_ID}."}, + ) + ) + + with ( + _mock_search_tool_backend(_scoping_db_tools()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + ui_team_is_not_a_real_team, + ), + _override_auth(dashboard_session_user), + ): + response = TestClient(app).get("/search_tools/list") + + assert response.status_code == 200 + names = {t["search_tool_name"] for t in response.json()["search_tools"]} + assert names == {"db-tool-1", "db-tool-2", "db-tool-3"} + ui_team_is_not_a_real_team.assert_not_awaited() + + +@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. + """ + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + restricted_dashboard_session = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=UI_SESSION_TOKEN_TEAM_ID, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + search_tools=["db-tool-3"], + ), + ) + lookup, asked_for = _recording_team_lookup() + + visible = await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), + restricted_dashboard_session, + lookup, + ) + + assert [t["search_tool_name"] for t in visible] == ["db-tool-3"] + assert asked_for == [] + + +@pytest.mark.asyncio +async def test_filter_visible_search_tools_still_applies_a_real_team_allowlist(): + """A caller with a real team is still resolved and scoped by that team's allowlist.""" + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="team-1", + ) + lookup, asked_for = _recording_team_lookup( + team_object=LiteLLM_TeamTable( + team_id="team-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-team", + search_tools=["db-tool-2"], + ), + ) + ) + + visible = await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), + team_member, + lookup, + ) + + assert [t["search_tool_name"] for t in visible] == ["db-tool-2"] + assert asked_for == ["team-1"] + + +@pytest.mark.asyncio +async def test_filter_visible_search_tools_propagates_a_real_team_lookup_failure(): + """ + A caller whose real team cannot be resolved must not fall through to "no team", which + would drop that team's allowlist and show tools the caller may not call. + """ + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="deleted-team", + ) + lookup, asked_for = _recording_team_lookup( + raises=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."}) + ) + + with pytest.raises(HTTPException) as exc_info: + await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2"), + team_member, + lookup, + ) + + assert exc_info.value.status_code == 404 + assert asked_for == ["deleted-team"] + + +@pytest.mark.asyncio +async def test_list_search_tools_reports_a_missing_real_team_as_404(): + """ + The endpoint surfaces a genuine team lookup failure with its own status instead of + masking it as a 500 or quietly returning an unscoped list. + """ + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="deleted-team", + ) + + with ( + _mock_search_tool_backend(_scoping_db_tools()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + AsyncMock( + side_effect=HTTPException( + status_code=404, + detail={"error": "Team doesn't exist in db. Team=deleted-team."}, + ) + ), + ), + _override_auth(team_member), + ): + response = TestClient(app).get("/search_tools/list") + + assert response.status_code == 404 + assert "search_tools" not in response.json()