scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(proxy): scope a dashboard session on /search_tools/list to its user's real teams

Browse source 
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.
This commit is contained in:
Yuneng Jiang 2026-08-06 12:07:09 -07:00
parent fe772cc0a1
commit cacf3debac
No known key found for this signature in database
3 changed files with 156 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/auth/auth_checks.py
View file

@ -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,

82
litellm/proxy/search_endpoints/search_tool_management.py
View file

@ -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(

153
tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py
View file

@ -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