From 0485a1859a94520ff55148aeab313a102d2b50be Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 17 Mar 2026 21:19:42 -0700 Subject: [PATCH] fix: use get_user_object helper, preserve caller org_id filter - Replace raw find_unique with get_user_object in _build_team_list_where_conditions for cache/metrics consistency - Remove over-complex OR clause for org admin + user_id: when user_id is provided, filter by that user's direct team memberships (same as regular users) since the access control gate already verified the org admin's authority - Preserve caller-supplied organization_id instead of overwriting with org_admin_org_ids - Update test mock to match get_user_object call path Co-Authored-By: Claude Opus 4.6 (1M context) --- .../management_endpoints/team_endpoints.py | 41 +++++++------- .../test_team_endpoints.py | 55 +++++++++++-------- 2 files changed, 53 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b32206a454..531cd67589 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3249,6 +3249,8 @@ async def _build_team_list_where_conditions( user_id: Optional[str], use_deleted_table: bool, org_admin_org_ids: Optional[List[str]] = None, + user_api_key_cache: Optional[Any] = None, + proxy_logging_obj: Optional[Any] = None, ) -> Optional[Dict[str, Any]]: """ Build where conditions for team list query. @@ -3274,34 +3276,33 @@ async def _build_team_list_where_conditions( where_conditions["organization_id"] = {"in": org_admin_org_ids} if user_id: - user_object = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} - ) - if user_object is None: + try: + user_object_correct_type = 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 user_object_correct_type is None: raise HTTPException( status_code=404, detail={"error": f"User not found, passed user_id={user_id}"}, ) - user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) user_team_ids = user_object_correct_type.teams or [] if use_deleted_table: where_conditions["members"] = {"has": user_id} - elif org_admin_org_ids is not None: - # Org admin with user_id filter: show teams in their orgs - # OR teams the user is a direct member of (matches legacy - # _authorize_and_filter_teams behaviour). - # When team_id is also provided, the exact match is already in - # where_conditions and the org scope just needs to be added — - # no OR expansion needed. - if team_id is not None: - where_conditions["organization_id"] = {"in": org_admin_org_ids} - elif user_team_ids: - org_condition: Dict[str, Any] = {"organization_id": {"in": org_admin_org_ids}} - where_conditions["OR"] = [org_condition, {"team_id": {"in": user_team_ids}}] - else: - where_conditions["organization_id"] = {"in": org_admin_org_ids} else: + # When user_id is provided, filter by that user's direct team + # memberships. For org admins the access control gate in + # list_team_v2 already verified the caller's authority — the + # filter logic is the same as for regular users. if not user_team_ids: return None # no memberships — skip the DB query elif team_id is not None: @@ -3466,6 +3467,8 @@ async def list_team_v2( user_id=user_id, use_deleted_table=use_deleted_table, org_admin_org_ids=org_admin_org_ids, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) if where_conditions is None: diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 41a724c271..122ac749f2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2142,19 +2142,21 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: + 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"): # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db - - # Mock user lookup - mock_user_object = Mock() - mock_user_object.model_dump.return_value = { - "user_id": "non_admin_user_123", - "teams": ["team_1", "team_2"], - } - mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user_object) - + + # Mock get_user_object to return a user with teams + from litellm.proxy._types import LiteLLM_UserTable + + mock_user = LiteLLM_UserTable( + user_id="non_admin_user_123", + teams=["team_1", "team_2"], + ) + # Mock team lookup mock_teams = [ Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Team 1"}), @@ -2163,21 +2165,26 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) - # Should NOT raise an exception - result = await list_team_v2( - http_request=mock_request, - user_id="non_admin_user_123", # Non-admin querying their own teams - user_api_key_dict=mock_user_api_key_dict_non_admin, - team_id=None, - page=1, - page_size=10, - status=None, - ) + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ): + # Should NOT raise an exception + result = await list_team_v2( + http_request=mock_request, + user_id="non_admin_user_123", # Non-admin querying their own teams + user_api_key_dict=mock_user_api_key_dict_non_admin, + team_id=None, + page=1, + page_size=10, + status=None, + ) - # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + # Should return results without error + assert "teams" in result + assert "total" in result + assert result["total"] == 2 @pytest.mark.asyncio