Merge pull request #22722 from atapia27/feat/org-exclusive-add-member

org-exclusive-add-member

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -30,6 +30,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
     get_daily_activity,
     get_daily_activity_aggregated,
 )
+from litellm.proxy.auth.auth_checks import get_user_object
 from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
 from litellm.proxy.management_endpoints.key_management_endpoints import (
     generate_key_helper_fn,
@@ -1844,7 +1845,11 @@ async def ui_view_users(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
-    [PROXY-ADMIN ONLY]Filter users based on partial match of user_id or email with pagination.
+    Filter users based on partial match of user_id or email with pagination.
+
+    - Proxy admins: receive all matching users.
+    - Organization admins: receive only users in their own organization(s).
+    - Other roles: access denied (403).
 
     Args:
         user_id (Optional[str]): Partial user ID to search for
@@ -1854,19 +1859,69 @@ async def ui_view_users(
         user_api_key_dict (UserAPIKeyAuth): User authentication information
 
     Returns:
-        List[LiteLLM_SpendLogs]: Paginated list of matching user records
+        List of matching user records (LiteLLM_UserTableFiltered), scoped by org for org admins.
     """
-    from litellm.proxy.proxy_server import prisma_client
+    from litellm.proxy.proxy_server import (
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
 
     if prisma_client is None:
         raise HTTPException(status_code=500, detail={"error": "No db connected"})
 
     try:
+        # Restrict by caller role: proxy admin sees all; org admin sees only their org(s); others 403
+        is_proxy_admin = _user_has_admin_view(user_api_key_dict)
+        if not is_proxy_admin:
+            if user_api_key_dict.user_id is None:
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "error": "Only proxy admins and organization admins can search users."
+                    },
+                )
+            try:
+                caller_user = 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,
+                    proxy_logging_obj=proxy_logging_obj,
+                )
+            except ValueError:
+                # get_user_object raises ValueError when user not found (user_id_upsert=False)
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "error": "Only proxy admins and organization admins can search users."
+                    },
+                )
+            if caller_user is None:
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "error": "Only proxy admins and organization admins can search users."
+                    },
+                )
+            org_admin_org_ids = [
+                m.organization_id
+                for m in (caller_user.organization_memberships or [])
+                if m.user_role == LitellmUserRoles.ORG_ADMIN.value
+            ]
+            if not org_admin_org_ids:
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "error": "Only proxy admins and organization admins can search users."
+                    },
+                )
+
         # Calculate offset for pagination
         skip = (page - 1) * page_size
 
         # Build where conditions based on provided parameters
-        where_conditions = {}
+        where_conditions: Dict[str, Any] = {}
 
         if user_id:
             where_conditions["user_id"] = {
@@ -1880,6 +1935,12 @@ async def ui_view_users(
                 "mode": "insensitive",  # Case-insensitive search
             }
 
+        # Org admins: only users in their org(s)
+        if not is_proxy_admin:
+            where_conditions["organization_memberships"] = {
+                "some": {"organization_id": {"in": org_admin_org_ids}}
+            }
+
         # Query users with pagination and filters
         users: Optional[List[BaseModel]] = (
             await prisma_client.db.litellm_usertable.find_many(
@@ -1895,6 +1956,8 @@ async def ui_view_users(
 
         return [LiteLLM_UserTableFiltered(**user.model_dump()) for user in users]
 
+    except HTTPException:
+        raise
     except Exception as e:
         verbose_proxy_logger.exception(f"Error searching users: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error searching users: {str(e)}")

tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -34,7 +34,8 @@ client = TestClient(app)
 @pytest.mark.asyncio
 async def test_ui_view_users_with_null_email(mocker, caplog):
     """
-    Test that /user/filter/ui endpoint returns users even when they have null email fields
+    Test that /user/filter/ui endpoint returns users even when they have null email fields.
+    Uses proxy admin so no org filtering is applied.
     """
     # Mock the prisma client
     mock_prisma_client = mocker.MagicMock()
@@ -48,19 +49,18 @@ async def test_ui_view_users_with_null_email(mocker, caplog):
         "created_at": "2024-01-01T00:00:00Z",
     }
 
-    # Setup the mock find_many response
-    # Setup the mock find_many response as an async function
     async def mock_find_many(*args, **kwargs):
         return [mock_user]
 
     mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
 
-    # Patch the prisma client import in the endpoint
     mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
 
-    # Call ui_view_users function directly
+    # Proxy admin: no org filter, no get_user_object call
     response = await ui_view_users(
-        user_api_key_dict=UserAPIKeyAuth(user_id="test_user"),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN
+        ),
         user_id="test_user",
         user_email=None,
         page=1,
@@ -72,6 +72,121 @@ async def test_ui_view_users_with_null_email(mocker, caplog):
     ]
 
 
+@pytest.mark.asyncio
+async def test_ui_view_users_proxy_admin_no_org_filter(mocker):
+    """
+    Proxy admin: find_many is called without organization_memberships in where.
+    """
+    mock_prisma_client = mocker.MagicMock()
+    async def mock_find_many(*args, **kwargs):
+        assert "organization_memberships" not in (kwargs.get("where") or {})
+        return []
+
+    mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
+    mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+
+    await ui_view_users(
+        user_api_key_dict=UserAPIKeyAuth(
+            user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
+        ),
+        user_id=None,
+        user_email="foo",
+        page=1,
+        page_size=50,
+    )
+
+
+@pytest.mark.asyncio
+async def test_ui_view_users_org_admin_filtered_by_org(mocker):
+    """
+    Org admin: find_many is called with organization_memberships filter so only users
+    in the caller's org(s) are returned.
+    """
+    from litellm.proxy._types import LiteLLM_OrganizationMembershipTable
+
+    mock_prisma_client = mocker.MagicMock()
+    org_id = "org-123"
+
+    async def mock_find_many(*args, **kwargs):
+        where = kwargs.get("where") or {}
+        assert "organization_memberships" in where
+        assert where["organization_memberships"] == {
+            "some": {"organization_id": {"in": [org_id]}}
+        }
+        return []
+
+    mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
+    mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
+    mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mocker.MagicMock())
+
+    caller_user = mocker.MagicMock()
+    caller_user.organization_memberships = [
+        LiteLLM_OrganizationMembershipTable(
+            user_id="org-admin",
+            organization_id=org_id,
+            user_role=LitellmUserRoles.ORG_ADMIN.value,
+            created_at=datetime.now(timezone.utc),
+            updated_at=datetime.now(timezone.utc),
+        )
+    ]
+
+    async def mock_get_user_object(*args, **kwargs):
+        return caller_user
+
+    mocker.patch(
+        "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_object",
+        side_effect=mock_get_user_object,
+    )
+
+    response = await ui_view_users(
+        user_api_key_dict=UserAPIKeyAuth(user_id="org-admin", user_role=None),
+        user_id=None,
+        user_email="u",
+        page=1,
+        page_size=50,
+    )
+
+    assert response == []
+
+
+@pytest.mark.asyncio
+async def test_ui_view_users_non_org_admin_returns_403(mocker):
+    """
+    Caller is not proxy admin and not org admin: endpoint returns 403.
+    """
+    from fastapi import HTTPException
+
+    mock_prisma_client = mocker.MagicMock()
+    mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock())
+    mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mocker.MagicMock())
+
+    # Caller has no org admin membership
+    caller_user = mocker.MagicMock()
+    caller_user.organization_memberships = []  # not an org admin
+
+    async def mock_get_user_object(*args, **kwargs):
+        return caller_user
+
+    mocker.patch(
+        "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_object",
+        side_effect=mock_get_user_object,
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await ui_view_users(
+            user_api_key_dict=UserAPIKeyAuth(user_id="internal_user", user_role=None),
+            user_id=None,
+            user_email="u",
+            page=1,
+            page_size=50,
+        )
+
+    assert exc_info.value.status_code == 403
+    assert "Only proxy admins and organization admins" in str(exc_info.value.detail)
+
+
 def test_user_daily_activity_types():
     """
     Assert all fiels in SpendMetrics are reported in DailySpendMetadata as "total_"