From 861bf37d8ceab42ed637089c4f73ae2b52cc06cd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 00:06:10 -0700 Subject: [PATCH] fix(proxy): stop resolving the UI session sentinel team on POST /search Every Admin UI session key is stamped with UI_SESSION_TOKEN_TEAM_ID ("litellm-dashboard"), a reserved sentinel that never has a row in LiteLLM_TeamTable because /team/new refuses to create it. The /search handler looked that id up as a real team, so get_team_object raised 404 and no non-admin could invoke a search tool from the dashboard. Resolve the allowlist team through a shared helper that returns None for the sentinel, so those callers stay scoped by their key-level allowlist alone. A real team id is still looked up and its allowlist still applies, and a lookup that genuinely fails still surfaces instead of falling through to "no team". --- litellm/proxy/search_endpoints/endpoints.py | 26 +- .../search_endpoints/search_tool_access.py | 47 ++++ .../proxy/search_endpoints/test_endpoints.py | 223 ++++++++++++++++++ 3 files changed, 277 insertions(+), 19 deletions(-) create mode 100644 litellm/proxy/search_endpoints/search_tool_access.py create mode 100644 tests/test_litellm/proxy/search_endpoints/test_endpoints.py diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 2676682c59d..2959abfa9bd 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -142,7 +142,9 @@ async def search( from litellm.proxy.auth.auth_checks import ( can_key_call_search_tool, can_team_call_search_tool, - get_team_object, + ) + from litellm.proxy.search_endpoints.search_tool_access import ( + resolve_allowlist_team, ) try: @@ -153,24 +155,10 @@ async def search( ) # Check team-level access if key is associated with a team - if user_api_key_dict.team_id: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - team_object: Final = 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, - ) - await can_team_call_search_tool( - search_tool_name=search_tool_name_value, - team_object=team_object, - ) + await can_team_call_search_tool( + search_tool_name=search_tool_name_value, + team_object=await resolve_allowlist_team(user_api_key_dict), + ) except Exception as e: verbose_proxy_logger.error("Search tool authorization failed for %s: %s", search_tool_name_value, e) raise diff --git a/litellm/proxy/search_endpoints/search_tool_access.py b/litellm/proxy/search_endpoints/search_tool_access.py new file mode 100644 index 00000000000..78e27915609 --- /dev/null +++ b/litellm/proxy/search_endpoints/search_tool_access.py @@ -0,0 +1,47 @@ +""" +Shared team scoping for the search tool authorization checks. +""" + +from collections.abc import Awaitable, Callable +from typing import Final, TypeAlias + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth + +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, + ) + + +async def resolve_allowlist_team( + user_api_key_dict: UserAPIKeyAuth, + lookup_team_object: TeamObjectLookup = team_object_from_db, +) -> LiteLLM_TeamTable | 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 await lookup_team_object(team_id, user_api_key_dict) diff --git a/tests/test_litellm/proxy/search_endpoints/test_endpoints.py b/tests/test_litellm/proxy/search_endpoints/test_endpoints.py new file mode 100644 index 00000000000..de5b0d24ff7 --- /dev/null +++ b/tests/test_litellm/proxy/search_endpoints/test_endpoints.py @@ -0,0 +1,223 @@ +import contextlib +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.proxy_server import app +from litellm.proxy.search_endpoints.search_tool_access import resolve_allowlist_team + +SEARCH_TOOLS = [ + { + "search_tool_name": "db-tool-1", + "litellm_params": {"search_provider": "perplexity", "api_key": "pplx-secret-1"}, + }, + { + "search_tool_name": "db-tool-2", + "litellm_params": {"search_provider": "tavily", "api_key": "tvly-secret-2"}, + }, + { + "search_tool_name": "db-tool-3", + "litellm_params": {"search_provider": "exa", "api_key": "exa-secret-3"}, + }, +] + + +@contextlib.contextmanager +def _override_auth(user): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: user + try: + yield + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + +@contextlib.contextmanager +def _mock_search_backend(lookup_team_object): + """Patch the router, prisma client, and downstream request processing so /search reaches + the authorization checks and, if they pass, returns without calling a provider.""" + router = MagicMock() + router.search_tools = SEARCH_TOOLS + processor = MagicMock() + processor.return_value.base_process_llm_request = AsyncMock(return_value={"object": "search", "results": []}) + with ( + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.auth.auth_checks.get_team_object", lookup_team_object), + patch( + "litellm.proxy.search_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + processor, + ), + ): + yield + + +def _team_ids_looked_up(lookup: AsyncMock) -> list[str]: + return [awaited.kwargs["team_id"] for awaited in lookup.await_args_list] + + +def _dashboard_session_key(search_tools: list[str] | None) -> UserAPIKeyAuth: + return 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=search_tools) + if search_tools is not None + else None + ), + ) + + +def _team_lookup_fails_with_404(team_id: str) -> AsyncMock: + return AsyncMock( + side_effect=HTTPException( + status_code=404, + detail={"error": f"Team doesn't exist in db. Team={team_id}."}, + ) + ) + + +@pytest.mark.asyncio +async def test_search_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, so no non-admin could invoke a search tool from the dashboard. + """ + ui_team_is_not_a_real_team = _team_lookup_fails_with_404(UI_SESSION_TOKEN_TEAM_ID) + + with ( + _mock_search_backend(ui_team_is_not_a_real_team), + _override_auth(_dashboard_session_key(["db-tool-3"])), + ): + response = TestClient(app).post("/search/db-tool-3", json={"query": "litellm"}) + + assert response.status_code == 200 + ui_team_is_not_a_real_team.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_search_dashboard_session_key_still_bound_by_its_key_allowlist(): + """ + Skipping the synthetic team must not widen access: a dashboard session whose key + allowlists only one tool is still refused every other tool. + """ + ui_team_is_not_a_real_team = _team_lookup_fails_with_404(UI_SESSION_TOKEN_TEAM_ID) + + with ( + _mock_search_backend(ui_team_is_not_a_real_team), + _override_auth(_dashboard_session_key(["db-tool-3"])), + ): + response = TestClient(app).post("/search/db-tool-1", json={"query": "litellm"}) + + assert response.status_code == 403 + assert "db-tool-1" in response.text + ui_team_is_not_a_real_team.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_search_real_team_allowlist_still_blocks_a_tool_it_does_not_permit(): + """A caller with a real team is still resolved and scoped by that team's allowlist.""" + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + 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"], + ), + ) + ) + + with _mock_search_backend(lookup), _override_auth(team_member): + blocked = TestClient(app).post("/search/db-tool-1", json={"query": "litellm"}) + allowed = TestClient(app).post("/search/db-tool-2", json={"query": "litellm"}) + + assert blocked.status_code == 403 + assert allowed.status_code == 200 + assert _team_ids_looked_up(lookup) == ["team-1", "team-1"] + + +@pytest.mark.asyncio +async def test_search_missing_real_team_is_still_rejected(): + """ + A caller whose real team cannot be resolved must not fall through to "no team", which + would drop that team's allowlist and let it call tools it may not. + """ + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="deleted-team", + ) + lookup = _team_lookup_fails_with_404("deleted-team") + + with _mock_search_backend(lookup), _override_auth(team_member): + response = TestClient(app).post("/search/db-tool-1", json={"query": "litellm"}) + + assert response.status_code == 404 + assert _team_ids_looked_up(lookup) == ["deleted-team"] + + +@pytest.mark.asyncio +async def test_resolve_allowlist_team_skips_only_the_ui_session_sentinel(): + lookup = AsyncMock() + + assert await resolve_allowlist_team(_dashboard_session_key(None), lookup) is None + lookup.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_allowlist_team_returns_none_when_the_key_has_no_team(): + keyless_of_team = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user") + lookup = AsyncMock() + + assert await resolve_allowlist_team(keyless_of_team, lookup) is None + lookup.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_allowlist_team_looks_up_a_real_team(): + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="team-1", + ) + team = LiteLLM_TeamTable(team_id="team-1") + lookup = AsyncMock(return_value=team) + + assert await resolve_allowlist_team(team_member, lookup) is team + assert [awaited.args[0] for awaited in lookup.await_args_list] == ["team-1"] + + +@pytest.mark.asyncio +async def test_resolve_allowlist_team_propagates_a_real_team_lookup_failure(): + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="deleted-team", + ) + lookup = _team_lookup_fails_with_404("deleted-team") + + with pytest.raises(HTTPException) as exc_info: + await resolve_allowlist_team(team_member, lookup) + + assert exc_info.value.status_code == 404