tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-20 00:18:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #20987 from BerriAI/litellm_inv_user_org

Browse Source 
[Feature] Allow Organization and Team Admins to call /invitation/new
This commit is contained in:
yuneng-jiang
2026-02-11 11:38:15 -08:00
committed by GitHub
parent 9975a9e3d4 40295595c7
commit b88b2520c0
5 changed files with 659 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+3
View File
@@ -632,6 +632,9 @@ class LiteLLMRoutes(enum.Enum):
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
# Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
"/invitation/new",
"/invitation/delete",
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
litellm/proxy/management_endpoints/common_utils.py
+149
View File
@@ -8,6 +8,7 @@ from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
LiteLLM_OrganizationTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
UserAPIKeyAuth,
)
@@ -108,6 +109,154 @@ async def _user_has_admin_privileges(
return False
def _org_admin_can_invite_user(
admin_user_obj: LiteLLM_UserTable,
target_user_obj: LiteLLM_UserTable,
) -> bool:
"""
Check if an org admin can invite the target user.
Target user must be in at least one org where the admin has org admin role.
Args:
admin_user_obj: The admin user's full object (from get_user_object)
target_user_obj: The target user's full object (from get_user_object)
Returns:
True if target user is in an org where admin has org admin role
"""
if admin_user_obj.organization_memberships is None:
return False
admin_org_ids = {
m.organization_id
for m in admin_user_obj.organization_memberships
if m.user_role == LitellmUserRoles.ORG_ADMIN.value
}
if not admin_org_ids:
return False
if target_user_obj.organization_memberships is None:
return False
target_org_ids = {
m.organization_id for m in target_user_obj.organization_memberships
}
return bool(admin_org_ids & target_org_ids)
async def _team_admin_can_invite_user(
user_api_key_dict: UserAPIKeyAuth,
admin_user_obj: LiteLLM_UserTable,
target_user_obj: LiteLLM_UserTable,
prisma_client: "PrismaClient",
) -> bool:
"""
Check if a team admin can invite the target user.
Target user must be in at least one team where the admin has team admin role.
Args:
user_api_key_dict: The admin user's API key auth object
admin_user_obj: The admin user's full object (from get_user_object)
target_user_obj: The target user's full object (from get_user_object)
prisma_client: Prisma client for database operations
Returns:
True if target user is in a team where admin has team admin role
"""
if not admin_user_obj.teams or len(admin_user_obj.teams) == 0:
return False
if not target_user_obj.teams or len(target_user_obj.teams) == 0:
return False
teams = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": admin_user_obj.teams}}
)
admin_team_ids = [
team.team_id
for team in teams
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict,
team_obj=LiteLLM_TeamTable(**team.model_dump()),
)
]
if not admin_team_ids:
return False
target_team_ids = set(target_user_obj.teams)
return bool(set(admin_team_ids) & target_team_ids)
async def admin_can_invite_user(
target_user_id: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Optional["PrismaClient"] = None,
user_api_key_cache: Optional["DualCache"] = None,
proxy_logging_obj: Optional["ProxyLogging"] = None,
) -> bool:
"""
Check if the admin can create an invitation for the target user.
- Proxy admins: can invite any user
- Org admins: can only invite users in their org(s)
- Team admins: can only invite users in their team(s)
Uses get_user_object for caching of both admin and target user objects.
Args:
target_user_id: The user_id of the user to invite
user_api_key_dict: The admin user's API key auth object
prisma_client: Prisma client for database operations
user_api_key_cache: Cache for user API keys
proxy_logging_obj: Proxy logging object
Returns:
True if user can invite the target user
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True
if prisma_client is None or user_api_key_dict.user_id is None:
return False
from litellm.caching import DualCache as DualCacheImport
from litellm.proxy.auth.auth_checks import get_user_object
try:
cache = user_api_key_cache or DualCacheImport()
admin_user_obj = await get_user_object(
user_id=user_api_key_dict.user_id,
prisma_client=prisma_client,
user_api_key_cache=cache,
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
)
if admin_user_obj is None:
return False
target_user_obj = await get_user_object(
user_id=target_user_id,
prisma_client=prisma_client,
user_api_key_cache=cache,
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
)
if target_user_obj is None:
return False
if _org_admin_can_invite_user(admin_user_obj, target_user_obj):
return True
if await _team_admin_can_invite_user(
user_api_key_dict=user_api_key_dict,
admin_user_obj=admin_user_obj,
target_user_obj=target_user_obj,
prisma_client=prisma_client,
):
return True
return False
except Exception as e:
verbose_proxy_logger.debug(
f"Error checking invite permission for user {user_api_key_dict.user_id}: {e}"
)
return False
def _set_object_metadata_field(
object_data: Union[
LiteLLM_TeamTable,
litellm/proxy/proxy_server.py
+60 -3
View File
@@ -338,7 +338,10 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import (
from litellm.proxy.management_endpoints.callback_management_endpoints import (
router as callback_management_endpoints_router,
)
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.proxy.management_endpoints.common_utils import (
admin_can_invite_user,
_user_has_admin_privileges,
)
from litellm.proxy.management_endpoints.cost_tracking_settings import (
router as cost_tracking_settings_router,
)
@@ -10380,7 +10383,17 @@ async def new_invitation(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
# Allow proxy admins and org/team admins (admin status from DB via get_user_object)
has_access = (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or await _user_has_admin_privileges(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
if not has_access:
raise HTTPException(
status_code=400,
detail={
@@ -10391,6 +10404,23 @@ async def new_invitation(
},
)
# Org/team admins can only invite users within their org/team
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
can_invite = await admin_can_invite_user(
target_user_id=data.user_id,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not can_invite:
raise HTTPException(
status_code=400,
detail={
"error": "You can only create invitations for users in your organization or team."
},
)
response = await create_invitation_for_user(
data=data,
user_api_key_dict=user_api_key_dict,
@@ -10543,7 +10573,16 @@ async def invitation_delete(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
# Proxy admins can delete any invitation; org admins only their own
is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
is_other_admin = await _user_has_admin_privileges(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not is_proxy_admin and not is_other_admin:
raise HTTPException(
status_code=400,
detail={
@@ -10554,6 +10593,24 @@ async def invitation_delete(
},
)
# Org admins can only delete invitations they created
if is_other_admin and not is_proxy_admin:
invitation = await prisma_client.db.litellm_invitationlink.find_unique(
where={"id": data.invitation_id}
)
if invitation is None:
raise HTTPException(
status_code=400,
detail={"error": "Invitation id does not exist in the database."},
)
if invitation.created_by != user_api_key_dict.user_id:
raise HTTPException(
status_code=403,
detail={
"error": "Organization admins can only delete invitations they created."
},
)
response = await prisma_client.db.litellm_invitationlink.delete(
where={"id": data.invitation_id}
)
tests/test_litellm/proxy/management_endpoints/test_common_utils.py
+327 -1
View File
@@ -7,10 +7,28 @@ enterprise (premium) license check, but should still be applied so that
users can intentionally clear previously-set fields.
"""
from unittest.mock import patch
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.proxy._types import (
Member,
LiteLLM_OrganizationMembershipTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin,
_org_admin_can_invite_user,
_set_object_metadata_field,
_team_admin_can_invite_user,
_update_metadata_fields,
_user_has_admin_privileges,
_user_has_admin_view,
admin_can_invite_user,
)
@@ -160,3 +178,311 @@ class TestUpdateMetadataFieldsEmptyCollections:
}
_update_metadata_fields(updated_kv=updated_kv)
mock_premium_check.assert_not_called()
class TestUserHasAdminView:
"""Tests for _user_has_admin_view function."""
@pytest.mark.parametrize(
"user_role,expected",
[
(LitellmUserRoles.PROXY_ADMIN, True),
(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, True),
(LitellmUserRoles.INTERNAL_USER, False),
(LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
],
)
def test_user_has_admin_view_by_role(self, user_role, expected):
"""Parametrized test: admin roles return True, non-admin return False."""
mock_auth = MagicMock()
mock_auth.user_role = user_role
assert _user_has_admin_view(mock_auth) == expected
def test_user_has_admin_view_with_user_api_key_auth(self):
"""Test with actual UserAPIKeyAuth object."""
auth_admin = UserAPIKeyAuth(
user_id="u1",
api_key="sk-xxx",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
auth_user = UserAPIKeyAuth(
user_id="u2",
api_key="sk-yyy",
user_role=LitellmUserRoles.INTERNAL_USER,
)
assert _user_has_admin_view(auth_admin) is True
assert _user_has_admin_view(auth_user) is False
class TestIsUserTeamAdmin:
"""Tests for _is_user_team_admin function."""
@pytest.mark.parametrize(
"members_with_roles,user_id,expected",
[
(
[Member(user_id="u1", role="admin")],
"u1",
True,
),
(
[Member(user_id="u1", role="user")],
"u1",
False,
),
(
[Member(user_id="u2", role="admin"), Member(user_id="u1", role="admin")],
"u1",
True,
),
([], "u1", False),
],
)
def test_is_user_team_admin_parametrized(
self, members_with_roles, user_id, expected
):
"""Parametrized test: user is team admin only when in members_with_roles with admin role."""
mock_auth = MagicMock()
mock_auth.user_id = user_id
team = LiteLLM_TeamTable(
team_id="team-1",
members_with_roles=members_with_roles,
)
assert _is_user_team_admin(mock_auth, team) == expected
def test_is_user_team_admin_user_not_in_team(self):
"""Test returns False when user is not in team members."""
auth = UserAPIKeyAuth(user_id="u99", api_key="sk-x", user_role=None)
team = LiteLLM_TeamTable(
team_id="team-1",
members_with_roles=[Member(user_id="u1", role="admin")],
)
assert _is_user_team_admin(auth, team) is False
class TestOrgAdminCanInviteUser:
"""Tests for _org_admin_can_invite_user function."""
def _make_membership(self, org_id: str, user_role: str):
now = datetime.now(timezone.utc)
return LiteLLM_OrganizationMembershipTable(
user_id="u",
organization_id=org_id,
user_role=user_role,
created_at=now,
updated_at=now,
)
@pytest.mark.parametrize(
"admin_orgs,target_orgs,expected",
[
(["org1"], ["org1"], True),
(["org1", "org2"], ["org2"], True),
(["org1"], ["org2"], False),
([], ["org1"], False),
(["org1"], [], False),
],
)
def test_org_admin_can_invite_user_parametrized(
self, admin_orgs, target_orgs, expected
):
"""Parametrized test: can invite when target is in org where admin has ORG_ADMIN role."""
admin_user = LiteLLM_UserTable(
user_id="admin",
organization_memberships=[
self._make_membership(oid, LitellmUserRoles.ORG_ADMIN.value)
for oid in admin_orgs
],
)
target_user = LiteLLM_UserTable(
user_id="target",
organization_memberships=[
self._make_membership(oid, LitellmUserRoles.INTERNAL_USER.value)
for oid in target_orgs
],
)
assert _org_admin_can_invite_user(admin_user, target_user) == expected
def test_org_admin_can_invite_user_no_shared_org(self):
"""Test returns False when admin has no org admin role."""
admin_user = LiteLLM_UserTable(
user_id="admin",
organization_memberships=[
self._make_membership("org1", LitellmUserRoles.INTERNAL_USER.value),
],
)
target_user = LiteLLM_UserTable(
user_id="target",
organization_memberships=[
self._make_membership("org1", LitellmUserRoles.INTERNAL_USER.value),
],
)
assert _org_admin_can_invite_user(admin_user, target_user) is False
class TestTeamAdminCanInviteUser:
"""Tests for _team_admin_can_invite_user async function."""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"admin_teams,target_teams,user_is_admin_in,expected",
[
(["t1"], ["t1"], ["t1"], True),
(["t1", "t2"], ["t2"], ["t1", "t2"], True),
(["t1"], ["t2"], ["t1"], False),
],
)
async def test_team_admin_can_invite_user_parametrized(
self, admin_teams, target_teams, user_is_admin_in, expected
):
"""Parametrized test: can invite when target shares a team where user is admin."""
mock_prisma = MagicMock()
mock_auth = MagicMock()
mock_auth.user_id = "admin"
admin_user = LiteLLM_UserTable(user_id="admin", teams=admin_teams)
target_user = LiteLLM_UserTable(user_id="target", teams=target_teams)
def make_team(tid, is_admin):
m = (
[{"user_id": "admin", "role": "admin"}]
if is_admin
else []
)
obj = MagicMock()
obj.team_id = tid
obj.model_dump = lambda: {"team_id": tid, "members_with_roles": m}
return obj
teams = [
make_team(tid, tid in user_is_admin_in) for tid in admin_teams
]
mock_prisma.db.litellm_teamtable.find_many = AsyncMock(
return_value=teams
)
result = await _team_admin_can_invite_user(
user_api_key_dict=mock_auth,
admin_user_obj=admin_user,
target_user_obj=target_user,
prisma_client=mock_prisma,
)
assert result == expected
@pytest.mark.asyncio
async def test_team_admin_can_invite_user_no_shared_team(self):
"""Test returns False when admin and target share no team."""
mock_prisma = MagicMock()
mock_auth = MagicMock()
mock_auth.user_id = "admin"
admin_user = LiteLLM_UserTable(user_id="admin", teams=[])
target_user = LiteLLM_UserTable(user_id="target", teams=["t1"])
result = await _team_admin_can_invite_user(
user_api_key_dict=mock_auth,
admin_user_obj=admin_user,
target_user_obj=target_user,
prisma_client=mock_prisma,
)
assert result is False
class TestUserHasAdminPrivileges:
"""Tests for _user_has_admin_privileges async function."""
@pytest.mark.asyncio
async def test_proxy_admin_has_privileges(self):
"""Proxy admin always has admin privileges."""
auth = UserAPIKeyAuth(
user_id="admin",
api_key="sk-x",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
result = await _user_has_admin_privileges(
user_api_key_dict=auth,
prisma_client=None,
)
assert result is True
@pytest.mark.asyncio
async def test_non_admin_no_prisma_returns_false(self):
"""Non-admin with no prisma connection has no privileges."""
auth = UserAPIKeyAuth(
user_id="user1",
api_key="sk-x",
user_role=LitellmUserRoles.INTERNAL_USER,
)
result = await _user_has_admin_privileges(
user_api_key_dict=auth,
prisma_client=None,
)
assert result is False
class TestAdminCanInviteUser:
"""Tests for admin_can_invite_user async function."""
@pytest.mark.asyncio
async def test_proxy_admin_can_invite_any_user(self):
"""Proxy admin can invite any user regardless of org/team."""
auth = UserAPIKeyAuth(
user_id="admin",
api_key="sk-x",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
result = await admin_can_invite_user(
target_user_id="any-user",
user_api_key_dict=auth,
prisma_client=None,
)
assert result is True
@pytest.mark.asyncio
async def test_non_admin_cannot_invite_without_prisma(self):
"""Non-admin with no prisma cannot invite."""
auth = UserAPIKeyAuth(
user_id="user1",
api_key="sk-x",
user_role=LitellmUserRoles.INTERNAL_USER,
)
result = await admin_can_invite_user(
target_user_id="other-user",
user_api_key_dict=auth,
prisma_client=None,
)
assert result is False
class TestSetObjectMetadataField:
"""Tests for _set_object_metadata_field function."""
@pytest.mark.parametrize(
"field_name,value,should_call_premium",
[
("guardrails", ["g1"], True),
("model_rpm_limit", {"gpt-4": 10}, False),
],
)
def test_set_object_metadata_field_parametrized(
self, field_name, value, should_call_premium
):
"""Parametrized test: premium fields trigger _premium_user_check."""
team = LiteLLM_TeamTable(team_id="t1", metadata={})
with patch(
"litellm.proxy.management_endpoints.common_utils._premium_user_check"
) as mock_premium:
_set_object_metadata_field(team, field_name, value)
if should_call_premium:
mock_premium.assert_called_once()
else:
mock_premium.assert_not_called()
assert team.metadata[field_name] == value
def test_set_object_metadata_field_initializes_metadata_if_none(self):
"""Test initializes metadata dict when object has None."""
team = LiteLLM_TeamTable(team_id="t1", metadata=None)
with patch(
"litellm.proxy.management_endpoints.common_utils._premium_user_check"
):
_set_object_metadata_field(team, "model_rpm_limit", {"x": 1})
assert team.metadata == {"model_rpm_limit": {"x": 1}}
tests/test_litellm/proxy/test_proxy_server.py
+120
View File
@@ -3203,3 +3203,123 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch):
assert result["general_settings"]["nested"]["key1"] == "updated_value1"
assert result["general_settings"]["nested"]["key2"] == "value2"
assert result["general_settings"]["nested"]["key3"] == "value3"
class TestInvitationEndpoints:
"""Tests for /invitation/new and /invitation/delete endpoints."""
@pytest.fixture
def client_with_auth(self):
"""Create a test client with admin authentication."""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.proxy_server import cleanup_router_config_variables
cleanup_router_config_variables()
filepath = os.path.dirname(os.path.abspath(__file__))
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
asyncio.run(initialize(config=config_fp, debug=True))
mock_auth = MagicMock()
mock_auth.user_id = "admin-user-id"
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
mock_auth.api_key = "sk-test"
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
return TestClient(app)
@pytest.mark.parametrize(
"endpoint,payload,mock_return",
[
(
"/invitation/new",
{"user_id": "target-user-123"},
{
"id": "inv-123",
"user_id": "target-user-123",
"is_accepted": False,
"accepted_at": None,
"expires_at": "2025-02-18T00:00:00",
"created_at": "2025-02-11T00:00:00",
"created_by": "admin-user-id",
"updated_at": "2025-02-11T00:00:00",
"updated_by": "admin-user-id",
},
),
(
"/invitation/delete",
{"invitation_id": "inv-456"},
{
"id": "inv-456",
"user_id": "target-user-123",
"is_accepted": False,
"accepted_at": None,
"expires_at": "2025-02-18T00:00:00",
"created_at": "2025-02-11T00:00:00",
"created_by": "admin-user-id",
"updated_at": "2025-02-11T00:00:00",
"updated_by": "admin-user-id",
},
),
],
)
def test_invitation_endpoints_proxy_admin_success(
self, client_with_auth, endpoint, payload, mock_return
):
"""Proxy admin can successfully create and delete invitations."""
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_invitationlink = MagicMock()
if endpoint == "/invitation/new":
mock_create = AsyncMock(return_value=mock_return)
with patch(
"litellm.proxy.management_helpers.user_invitation.create_invitation_for_user",
mock_create,
):
response = client_with_auth.post(endpoint, json=payload)
else:
mock_prisma.db.litellm_invitationlink.find_unique = AsyncMock(
return_value={**mock_return, "created_by": "admin-user-id"}
)
mock_prisma.db.litellm_invitationlink.delete = AsyncMock(
return_value=mock_return
)
response = client_with_auth.post(endpoint, json=payload)
assert response.status_code == 200
data = response.json()
assert data["id"] == mock_return["id"]
assert data["user_id"] == mock_return["user_id"]
@pytest.mark.parametrize(
"endpoint,payload",
[
("/invitation/new", {"user_id": "target-user-123"}),
("/invitation/delete", {"invitation_id": "inv-456"}),
],
)
def test_invitation_endpoints_non_admin_denied(
self, client_with_auth, endpoint, payload
):
"""Non-admin users cannot access invitation endpoints."""
from litellm.proxy._types import LitellmUserRoles
mock_auth = MagicMock()
mock_auth.user_id = "regular-user"
mock_auth.user_role = LitellmUserRoles.INTERNAL_USER
mock_auth.api_key = "sk-regular"
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_invitationlink = MagicMock()
# Avoid triggering async DB calls in _user_has_admin_privileges
with patch(
"litellm.proxy.proxy_server._user_has_admin_privileges",
new_callable=AsyncMock,
return_value=False,
):
response = client_with_auth.post(endpoint, json=payload)
assert response.status_code == 400
body = response.json()
# ProxyException handler returns {"error": {...}}, HTTPException returns {"detail": {...}}
error_content = body.get("error", body.get("detail", body))
assert "not allowed" in str(error_content).lower()