mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): scope v1 /team/list teams via user.teams
Align legacy team listing with /v2/team/list for internal users by loading teams from LiteLLM_UserTable.teams instead of scanning members_with_roles. Auto-scope missing user_id to caller id. Adds mocked regression tests.
This commit is contained in:
Milan
parent
1d56e732e8
commit
687dd33d8a
2 changed files with 263 additions and 23 deletions
|
|
@ -4147,12 +4147,6 @@ async def _authorize_and_filter_teams(
|
|||
allowed_org_ids: Optional[List[str]] = None
|
||||
|
||||
if not is_proxy_admin:
|
||||
is_own_query = (
|
||||
user_id is not None
|
||||
and user_api_key_dict.user_id is not None
|
||||
and user_api_key_dict.user_id == user_id
|
||||
)
|
||||
|
||||
# Check if user is an org admin (even for own queries, so they see org teams)
|
||||
if user_api_key_dict.user_id is not None:
|
||||
caller_user = await get_user_object(
|
||||
|
|
@ -4172,15 +4166,23 @@ async def _authorize_and_filter_teams(
|
|||
if not allowed_org_ids:
|
||||
allowed_org_ids = None
|
||||
|
||||
if allowed_org_ids is None and not is_own_query:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "Only admin users can query all teams/other teams. Your user role={}".format(
|
||||
user_api_key_dict.user_role
|
||||
)
|
||||
},
|
||||
if allowed_org_ids is None:
|
||||
if user_id is None:
|
||||
user_id = user_api_key_dict.user_id
|
||||
is_own_query = (
|
||||
user_id is not None
|
||||
and user_api_key_dict.user_id is not None
|
||||
and user_api_key_dict.user_id == user_id
|
||||
)
|
||||
if not is_own_query:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "Only admin users can query all teams/other teams. Your user role={}".format(
|
||||
user_api_key_dict.user_role
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
if allowed_org_ids is not None:
|
||||
# Org admin: query DB for teams in their orgs
|
||||
|
|
@ -4198,16 +4200,36 @@ async def _authorize_and_filter_teams(
|
|||
and any(m.get("user_id") == user_id for m in team.members_with_roles)
|
||||
]
|
||||
elif user_id:
|
||||
# Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays)
|
||||
response = await prisma_client.db.litellm_teamtable.find_many(
|
||||
include={"litellm_model_table": True}
|
||||
# Regular users are scoped to the canonical team ids on their user row.
|
||||
# This mirrors /v2/team/list and avoids exposing teams from stale
|
||||
# members_with_roles entries.
|
||||
try:
|
||||
target_user = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except ValueError:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"User not found, passed user_id={user_id}"},
|
||||
)
|
||||
if target_user is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"User not found, passed user_id={user_id}"},
|
||||
)
|
||||
user_team_ids = target_user.teams or []
|
||||
if not user_team_ids:
|
||||
return []
|
||||
return list(
|
||||
await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"team_id": {"in": user_team_ids}},
|
||||
include={"litellm_model_table": True},
|
||||
)
|
||||
)
|
||||
return [
|
||||
team
|
||||
for team in response
|
||||
if team.members_with_roles
|
||||
and any(m.get("user_id") == user_id for m in team.members_with_roles)
|
||||
]
|
||||
else:
|
||||
# Proxy admin: all teams
|
||||
return list(
|
||||
|
|
|
|||
|
|
@ -2258,6 +2258,224 @@ async def test_bulk_team_member_add_no_db_connection():
|
|||
assert "DB not connected" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v1_internal_user_scoped_to_user_table_teams():
|
||||
"""
|
||||
Test that v1 /team/list scopes internal users to the canonical teams on
|
||||
LiteLLM_UserTable.teams, matching /v2/team/list behavior.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="internal_user_123",
|
||||
)
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="internal_user_123",
|
||||
teams=["allowed_team"],
|
||||
)
|
||||
|
||||
mock_team = Mock()
|
||||
mock_team.team_id = "allowed_team"
|
||||
mock_team.model_dump.return_value = {
|
||||
"team_id": "allowed_team",
|
||||
"team_alias": "Allowed Team",
|
||||
"members_with_roles": [
|
||||
{"user_id": "internal_user_123", "role": "user"},
|
||||
],
|
||||
}
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache"),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj"),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_user,
|
||||
),
|
||||
):
|
||||
mock_db = Mock()
|
||||
mock_prisma_client.db = mock_db
|
||||
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
|
||||
mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[])
|
||||
mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await list_team(
|
||||
http_request=mock_request,
|
||||
user_id="internal_user_123",
|
||||
organization_id=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].team_id == "allowed_team"
|
||||
team_query = mock_db.litellm_teamtable.find_many.call_args.kwargs
|
||||
assert team_query["where"] == {"team_id": {"in": ["allowed_team"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v1_internal_user_without_user_id_scopes_to_self():
|
||||
"""
|
||||
Test that v1 /team/list auto-scopes internal users without an explicit
|
||||
user_id to their own user_id instead of listing all teams.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="internal_user_123",
|
||||
)
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="internal_user_123",
|
||||
teams=["self_team"],
|
||||
)
|
||||
|
||||
mock_team = Mock()
|
||||
mock_team.team_id = "self_team"
|
||||
mock_team.model_dump.return_value = {
|
||||
"team_id": "self_team",
|
||||
"team_alias": "Self Team",
|
||||
"members_with_roles": [
|
||||
{"user_id": "internal_user_123", "role": "user"},
|
||||
],
|
||||
}
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache"),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj"),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_user,
|
||||
),
|
||||
):
|
||||
mock_db = Mock()
|
||||
mock_prisma_client.db = mock_db
|
||||
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
|
||||
mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[])
|
||||
mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await list_team(
|
||||
http_request=mock_request,
|
||||
user_id=None,
|
||||
organization_id=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].team_id == "self_team"
|
||||
team_query = mock_db.litellm_teamtable.find_many.call_args.kwargs
|
||||
assert team_query["where"] == {"team_id": {"in": ["self_team"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v1_internal_user_cannot_query_other_user():
|
||||
"""
|
||||
Test that v1 /team/list rejects internal users querying another user's teams.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="internal_user_123",
|
||||
)
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="internal_user_123",
|
||||
teams=["team_1"],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache"),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj"),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_user,
|
||||
),
|
||||
):
|
||||
mock_prisma_client.db = Mock()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await list_team(
|
||||
http_request=mock_request,
|
||||
user_id="other_user_456",
|
||||
organization_id=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert "Only admin users can query all teams/other teams" in str(
|
||||
exc_info.value.detail
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v1_proxy_admin_can_query_all_teams():
|
||||
"""
|
||||
Test that v1 /team/list keeps proxy admin behavior unchanged.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="admin_user_123",
|
||||
)
|
||||
|
||||
mock_team = Mock()
|
||||
mock_team.team_id = "team_1"
|
||||
mock_team.model_dump.return_value = {
|
||||
"team_id": "team_1",
|
||||
"team_alias": "Team 1",
|
||||
"members_with_roles": [],
|
||||
}
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client:
|
||||
mock_db = Mock()
|
||||
mock_prisma_client.db = mock_db
|
||||
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
|
||||
mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[])
|
||||
mock_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await list_team(
|
||||
http_request=mock_request,
|
||||
user_id=None,
|
||||
organization_id=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].team_id == "team_1"
|
||||
team_query = mock_db.litellm_teamtable.find_many.call_args.kwargs
|
||||
assert "where" not in team_query
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v2_security_check_non_admin_user():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue