diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 2959abfa9bd..7f195342dc8 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -139,25 +139,16 @@ async def search( search_tool_name_value: Final = data["search_tool_name"] # Authorization check: verify key can access this search tool - from litellm.proxy.auth.auth_checks import ( - can_key_call_search_tool, - can_team_call_search_tool, - ) from litellm.proxy.search_endpoints.search_tool_access import ( - resolve_allowlist_team, + authorize_search_tool, + resolve_allowlist_teams, ) try: - # Check key-level access - await can_key_call_search_tool( + await authorize_search_tool( search_tool_name=search_tool_name_value, - valid_token=user_api_key_dict, - ) - - # Check team-level access if key is associated with a team - await can_team_call_search_tool( - search_tool_name=search_tool_name_value, - team_object=await resolve_allowlist_team(user_api_key_dict), + user_api_key_dict=user_api_key_dict, + teams=await resolve_allowlist_teams(user_api_key_dict), ) except Exception as e: verbose_proxy_logger.error("Search tool authorization failed for %s: %s", search_tool_name_value, e) diff --git a/litellm/proxy/search_endpoints/search_tool_access.py b/litellm/proxy/search_endpoints/search_tool_access.py index 78e27915609..54776a0e836 100644 --- a/litellm/proxy/search_endpoints/search_tool_access.py +++ b/litellm/proxy/search_endpoints/search_tool_access.py @@ -2,13 +2,23 @@ Shared team scoping for the search tool authorization checks. """ -from collections.abc import Awaitable, Callable -from typing import Final, TypeAlias +import asyncio +from collections.abc import Sequence +from typing import Final, Literal, Protocol +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_TeamTable, ProxyException, UserAPIKeyAuth -TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]] + +class TeamObjectLookup(Protocol): + async def __call__(self, team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: ... + + +class SessionTeamIdsLookup(Protocol): + async def __call__(self, user_api_key_dict: UserAPIKeyAuth) -> Sequence[str]: ... async def team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: @@ -28,20 +38,142 @@ async def team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) - ) -async def resolve_allowlist_team( +async def session_team_ids_from_db(user_api_key_dict: UserAPIKeyAuth) -> Sequence[str]: + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + user_object: Final = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_object is None: + raise ValueError("Cannot resolve the teams of a dashboard session key that carries no user id.") + return tuple(user_object.teams or ()) + + +async def _load_session_team( + team_id: str, + user_api_key_dict: UserAPIKeyAuth, + lookup_team_object: TeamObjectLookup, +) -> LiteLLM_TeamTable | HTTPException: + try: + return await lookup_team_object(team_id, user_api_key_dict) + except HTTPException as exc: + verbose_proxy_logger.warning( + "Search tool scoping: dropping team %s for dashboard session user %s: %s", + team_id, + user_api_key_dict.user_id, + exc.detail, + ) + return exc + + +async def _session_teams( + user_api_key_dict: UserAPIKeyAuth, + lookup_team_object: TeamObjectLookup, + lookup_session_team_ids: SessionTeamIdsLookup, +) -> tuple[LiteLLM_TeamTable, ...]: + team_ids: Final = await lookup_session_team_ids(user_api_key_dict) + loaded: Final = await asyncio.gather( + *(_load_session_team(team_id, user_api_key_dict, lookup_team_object) for team_id in team_ids) + ) + teams: Final = tuple(entry for entry in loaded if isinstance(entry, LiteLLM_TeamTable)) + if loaded and not teams: + raise next(entry for entry in loaded if isinstance(entry, HTTPException)) + return teams + + +async def resolve_allowlist_teams( user_api_key_dict: UserAPIKeyAuth, lookup_team_object: TeamObjectLookup = team_object_from_db, -) -> LiteLLM_TeamTable | None: + lookup_session_team_ids: SessionTeamIdsLookup = session_team_ids_from_db, +) -> tuple[LiteLLM_TeamTable, ...]: """ - The team whose object_permission allowlist scopes this caller, or None when there is none. + The teams whose object_permission allowlists scope this caller. - 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. + A virtual key names exactly one team and that lookup still surfaces its failure, so a key + pointing at a team that does not exist is rejected rather than treated as unscoped. + + An Admin UI session key instead carries UI_SESSION_TOKEN_TEAM_ID, a reserved sentinel that + never has a row in LiteLLM_TeamTable (`/team/new` rejects it as a team id), so looking it up + would raise 404 for every dashboard caller. It resolves to the real teams backing the session + user, the same identity the MCP dashboard surfaces resolve, which keeps the team allowlists + binding on the dashboard instead of dropping them. + + Resolving that set never fails open. Loading the session user surfaces its own failure, and one + unloadable team is dropped only because `/team/delete` leaves its id behind on the user row, so + a stale membership is ordinary and dropping it can only narrow a union the other teams still + scope. Once no team survives there is no union left to narrow, so the first failure is raised + rather than handing back the empty set that means "belongs to no team". """ 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) + if not team_id: + return () + if team_id != UI_SESSION_TOKEN_TEAM_ID: + return (await lookup_team_object(team_id, user_api_key_dict),) + return await _session_teams(user_api_key_dict, lookup_team_object, lookup_session_team_ids) + + +async def _team_denial(search_tool_name: str, team: LiteLLM_TeamTable) -> ProxyException | None: + from litellm.proxy.auth.auth_checks import can_team_call_search_tool + + try: + await can_team_call_search_tool(search_tool_name=search_tool_name, team_object=team) + except ProxyException as denial: + return denial + return None + + +async def authorize_search_tool( + search_tool_name: str, + user_api_key_dict: UserAPIKeyAuth, + teams: Sequence[LiteLLM_TeamTable], +) -> Literal[True]: + """ + Enforce the key and team object_permission allowlists for one search tool. + + The caller is authorized when its key permits the tool and at least one of its teams does. + A caller with no teams is unrestricted at the team level, matching a key that belongs to no + team, and is still bound by its key allowlist. + """ + from litellm.proxy.auth.auth_checks import ( + can_key_call_search_tool, + can_team_call_search_tool, + ) + + await can_key_call_search_tool(search_tool_name=search_tool_name, valid_token=user_api_key_dict) + if not teams: + return await can_team_call_search_tool(search_tool_name=search_tool_name, team_object=None) + + denials: Final = await asyncio.gather(*(_team_denial(search_tool_name, team) for team in teams)) + if any(denial is None for denial in denials): + return True + raise next(denial for denial in denials if denial is not None) + + +async def can_view_search_tool( + search_tool_name: str, + user_api_key_dict: UserAPIKeyAuth, + teams: Sequence[LiteLLM_TeamTable], +) -> bool: + """ + Boolean variant of authorize_search_tool, so a listing surface shows exactly the tools the + caller may invoke. + """ + try: + await authorize_search_tool( + search_tool_name=search_tool_name, + user_api_key_dict=user_api_key_dict, + teams=teams, + ) + except ProxyException: + return False + return True diff --git a/tests/test_litellm/proxy/search_endpoints/test_endpoints.py b/tests/test_litellm/proxy/search_endpoints/test_endpoints.py index de5b0d24ff7..cdd9c6ba2ae 100644 --- a/tests/test_litellm/proxy/search_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/search_endpoints/test_endpoints.py @@ -13,11 +13,11 @@ from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth, ) from litellm.proxy.proxy_server import app -from litellm.proxy.search_endpoints.search_tool_access import resolve_allowlist_team SEARCH_TOOLS = [ { @@ -47,17 +47,19 @@ def _override_auth(user): @contextlib.contextmanager -def _mock_search_backend(lookup_team_object): +def _mock_search_backend(lookup_team_object, session_team_ids=()): """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": []}) + session_user = LiteLLM_UserTable(user_id="internal_user", teams=list(session_team_ids)) 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.auth.auth_checks.get_user_object", AsyncMock(return_value=session_user)), patch( "litellm.proxy.search_endpoints.endpoints.ProxyBaseLLMRequestProcessing", processor, @@ -83,6 +85,31 @@ def _dashboard_session_key(search_tools: list[str] | None) -> UserAPIKeyAuth: ) +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, + ), + ) + + +def _lookup_returning(*teams: LiteLLM_TeamTable) -> AsyncMock: + by_id = {team.team_id: team for team in teams} + + async def lookup(**kwargs): + team_id = kwargs["team_id"] + if team_id not in by_id: + raise HTTPException( + status_code=404, + detail={"error": f"Team doesn't exist in db. Team={team_id}."}, + ) + return by_id[team_id] + + return AsyncMock(side_effect=lookup) + + def _team_lookup_fails_with_404(team_id: str) -> AsyncMock: return AsyncMock( side_effect=HTTPException( @@ -92,12 +119,53 @@ def _team_lookup_fails_with_404(team_id: str) -> AsyncMock: ) +@pytest.mark.asyncio +async def test_search_dashboard_session_is_scoped_by_the_users_real_team_allowlist(): + """ + Regression: the Admin UI session key carries the reserved team id ``litellm-dashboard``, which + has no row in LiteLLM_TeamTable. Treating that as "no team" dropped team-level authorization + entirely, so a dashboard session with no key allowlist could invoke every tool on the proxy. + The session must be scoped by the allowlists of the teams its user actually belongs to. + """ + lookup = _lookup_returning(_team_allowing("team-1", ["db-tool-2"])) + + with ( + _mock_search_backend(lookup, session_team_ids=["team-1"]), + _override_auth(_dashboard_session_key(None)), + ): + 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 "db-tool-1" in blocked.text + assert allowed.status_code == 200 + assert _team_ids_looked_up(lookup) == ["team-1", "team-1"] + + +@pytest.mark.asyncio +async def test_search_dashboard_session_is_permitted_by_any_of_the_users_teams(): + """A user on several teams may invoke a tool any one of those teams allowlists.""" + lookup = _lookup_returning( + _team_allowing("team-1", ["db-tool-2"]), + _team_allowing("team-2", ["db-tool-3"]), + ) + + with ( + _mock_search_backend(lookup, session_team_ids=["team-1", "team-2"]), + _override_auth(_dashboard_session_key(None)), + ): + allowed = TestClient(app).post("/search/db-tool-3", json={"query": "litellm"}) + blocked = TestClient(app).post("/search/db-tool-1", json={"query": "litellm"}) + + assert allowed.status_code == 200 + assert blocked.status_code == 403 + + @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. + Regression: resolving ``litellm-dashboard`` 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) @@ -112,22 +180,40 @@ async def test_search_dashboard_session_key_does_not_look_up_the_ui_team(): @pytest.mark.asyncio -async def test_search_dashboard_session_key_still_bound_by_its_key_allowlist(): +async def test_search_dashboard_session_denied_when_none_of_the_users_teams_load(): """ - Skipping the synthetic team must not widen access: a dashboard session whose key - allowlists only one tool is still refused every other tool. + Regression: a dashboard session whose every team failed to load resolved to an empty team set, + which reads as "belongs to no team" and skips team authorization altogether. A user who could + invoke nothing must not become a user who can invoke everything. """ - ui_team_is_not_a_real_team = _team_lookup_fails_with_404(UI_SESSION_TOKEN_TEAM_ID) + every_team_is_gone = _team_lookup_fails_with_404("team-1") with ( - _mock_search_backend(ui_team_is_not_a_real_team), + _mock_search_backend(every_team_is_gone, session_team_ids=["team-1", "team-2"]), + _override_auth(_dashboard_session_key(None)), + ): + response = TestClient(app).post("/search/db-tool-1", json={"query": "litellm"}) + + assert response.status_code == 404 + assert every_team_is_gone.await_count == 2 + + +@pytest.mark.asyncio +async def test_search_dashboard_session_key_still_bound_by_its_key_allowlist(): + """ + Resolving the session through its user's teams must not widen access: a dashboard session + whose key allowlists only one tool is still refused every other tool. + """ + lookup = _lookup_returning(_team_allowing("team-1", [])) + + with ( + _mock_search_backend(lookup, session_team_ids=["team-1"]), _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 @@ -138,15 +224,7 @@ async def test_search_real_team_allowlist_still_blocks_a_tool_it_does_not_permit 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 = _lookup_returning(_team_allowing("team-1", ["db-tool-2"])) with _mock_search_backend(lookup), _override_auth(team_member): blocked = TestClient(app).post("/search/db-tool-1", json={"query": "litellm"}) @@ -175,49 +253,3 @@ async def test_search_missing_real_team_is_still_rejected(): 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 diff --git a/tests/test_litellm/proxy/search_endpoints/test_search_tool_access.py b/tests/test_litellm/proxy/search_endpoints/test_search_tool_access.py new file mode 100644 index 00000000000..e4dcabb5db8 --- /dev/null +++ b/tests/test_litellm/proxy/search_endpoints/test_search_tool_access.py @@ -0,0 +1,217 @@ +import os +import sys +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +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, + ProxyException, + UserAPIKeyAuth, +) +from litellm.proxy.search_endpoints.search_tool_access import ( + authorize_search_tool, + resolve_allowlist_teams, + session_team_ids_from_db, +) + + +def _team(team_id: str, search_tools: list[str] | None = None) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission=( + LiteLLM_ObjectPermissionTable(object_permission_id=f"op-{team_id}", search_tools=search_tools) + if search_tools is not None + else None + ), + ) + + +def _dashboard_session(search_tools: list[str] | None = 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 _key_on_team(team_id: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=team_id, + ) + + +def _lookup_failing_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_resolve_returns_no_teams_when_the_key_has_none(): + lookup = AsyncMock() + session_team_ids = AsyncMock() + + assert await resolve_allowlist_teams(_key_on_team(""), lookup, session_team_ids) == () + lookup.assert_not_awaited() + session_team_ids.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_looks_up_a_real_team_and_never_widens_to_the_session_path(): + team = _team("team-1") + lookup = AsyncMock(return_value=team) + session_team_ids = AsyncMock() + + assert await resolve_allowlist_teams(_key_on_team("team-1"), lookup, session_team_ids) == (team,) + assert [awaited.args[0] for awaited in lookup.await_args_list] == ["team-1"] + session_team_ids.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_propagates_a_real_team_lookup_failure(): + """A key naming a team that does not exist must be rejected, not treated as unscoped.""" + lookup = _lookup_failing_with_404("deleted-team") + + with pytest.raises(HTTPException) as exc_info: + await resolve_allowlist_teams(_key_on_team("deleted-team"), lookup, AsyncMock()) + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_resolve_maps_the_session_sentinel_onto_the_users_real_teams(): + """The sentinel itself is never looked up; the session user's own teams are.""" + teams = (_team("team-1"), _team("team-2")) + lookup = AsyncMock(side_effect=teams) + session_team_ids = AsyncMock(return_value=["team-1", "team-2"]) + + assert await resolve_allowlist_teams(_dashboard_session(), lookup, session_team_ids) == teams + assert [awaited.args[0] for awaited in lookup.await_args_list] == ["team-1", "team-2"] + + +@pytest.mark.asyncio +async def test_resolve_drops_only_the_session_team_that_cannot_be_loaded(): + """Dropping an unloadable team can only narrow the union, so the rest still bind.""" + loadable = _team("team-2") + + async def lookup_impl(team_id, _user_api_key_dict): + if team_id == "team-1": + raise HTTPException(status_code=404, detail={"error": "Team doesn't exist in db. Team=team-1."}) + return loadable + + session_team_ids = AsyncMock(return_value=["team-1", "team-2"]) + + resolved = await resolve_allowlist_teams(_dashboard_session(), AsyncMock(side_effect=lookup_impl), session_team_ids) + + assert resolved == (loadable,) + + +@pytest.mark.asyncio +async def test_resolve_denies_a_session_when_no_team_can_be_loaded(): + """ + The collapse case. Dropping teams narrows a union only while one survives; once none do there + is nothing left to narrow, and returning () would read as 'belongs to no team' and go + unrestricted. The failure has to surface instead. + """ + lookup = _lookup_failing_with_404("team-1") + session_team_ids = AsyncMock(return_value=["team-1", "team-2"]) + + with pytest.raises(HTTPException) as exc_info: + await resolve_allowlist_teams(_dashboard_session(), lookup, session_team_ids) + + assert exc_info.value.status_code == 404 + assert lookup.await_count == 2 + + +@pytest.mark.asyncio +async def test_resolve_surfaces_a_failure_to_resolve_the_session_users_teams(): + """A session user whose own record cannot be read is not a session user with no teams.""" + session_team_ids = AsyncMock(side_effect=ValueError("User doesn't exist in db. 'user_id'=internal_user.")) + + with pytest.raises(ValueError, match="User doesn't exist in db"): + await resolve_allowlist_teams(_dashboard_session(), AsyncMock(), session_team_ids) + + +@pytest.mark.asyncio +async def test_resolve_returns_no_teams_for_a_session_user_on_no_teams(): + """The one empty result that is genuinely unscoped, and the reason the collapse needs its own case.""" + lookup = AsyncMock() + + assert await resolve_allowlist_teams(_dashboard_session(), lookup, AsyncMock(return_value=[])) == () + lookup.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_session_team_ids_do_not_swallow_a_failure_to_load_the_session_user(monkeypatch): + """Resolution reads the session user directly so a DB failure cannot arrive as an empty team list.""" + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + with pytest.raises(Exception, match="No db connected"): + await session_team_ids_from_db(_dashboard_session()) + + +@pytest.mark.asyncio +async def test_authorize_permits_a_tool_any_one_team_allowlists(): + assert ( + await authorize_search_tool( + search_tool_name="db-tool-2", + user_api_key_dict=_dashboard_session(), + teams=(_team("team-1", ["db-tool-1"]), _team("team-2", ["db-tool-2"])), + ) + is True + ) + + +@pytest.mark.asyncio +async def test_authorize_denies_a_tool_no_team_allowlists(): + with pytest.raises(ProxyException) as exc_info: + await authorize_search_tool( + search_tool_name="db-tool-3", + user_api_key_dict=_dashboard_session(), + teams=(_team("team-1", ["db-tool-1"]), _team("team-2", ["db-tool-2"])), + ) + + assert exc_info.value.code == "403" + assert "db-tool-3" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_authorize_is_unrestricted_at_team_level_with_no_teams(): + assert ( + await authorize_search_tool( + search_tool_name="db-tool-3", + user_api_key_dict=_dashboard_session(), + teams=(), + ) + is True + ) + + +@pytest.mark.asyncio +async def test_authorize_enforces_the_key_allowlist_before_any_team_grant(): + """A team grant never widens a key past its own allowlist.""" + with pytest.raises(ProxyException) as exc_info: + await authorize_search_tool( + search_tool_name="db-tool-3", + user_api_key_dict=_dashboard_session(["db-tool-1"]), + teams=(_team("team-1", ["db-tool-3"]),), + ) + + assert exc_info.value.code == "403"