mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(proxy): scope a dashboard session on POST /search to its user's real teams
Treating the reserved litellm-dashboard team id as "no team" dropped team-level authorization for every Admin UI session. A session key carries no allowlist of its own, so both gates passed unconditionally and a non-admin could invoke any search tool on the proxy from the dashboard. The sentinel now resolves to the real teams backing the session user, the same identity the MCP dashboard surfaces resolve, and the caller is authorized when at least one of those teams allowlists the tool. A virtual key naming a real team is unchanged: one lookup, and a failed lookup still surfaces. Resolving that set never falls back to the empty one. An unloadable team is dropped only because /team/delete leaves its id on the user row, so a stale membership is ordinary and dropping it narrows a union the other teams still scope. Once no team survives there is nothing left to narrow, so the failure is raised rather than handed back as the empty set that means "belongs to no team". Session teams are read from the user record directly for the same reason: the MCP helper answers a failed read with an empty list, which is indistinguishable from a user who genuinely belongs to no team.
This commit is contained in:
Yuneng Jiang
parent
23cb16a8d2
commit
fe772cc0a1
4 changed files with 467 additions and 95 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue