diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6b4cf72fc1..125efb03d8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -62,11 +62,9 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - _get_mcp_server_ids_from_access_groups, - _get_models_from_access_groups, allowed_route_check_inside_route, can_org_access_model, + get_access_object, get_org_object, get_team_object, get_user_object, @@ -3349,6 +3347,9 @@ async def _resolve_access_group_resources( """ Resolve resources inherited from access groups. + Fetches each access group object once and extracts all three resource + fields in a single pass (models, MCP servers, agents). + Returns only the access-group-sourced resources (not direct assignments). Keeps them separate so callers can distinguish where each resource comes from. """ @@ -3360,16 +3361,38 @@ async def _resolve_access_group_resources( if not access_group_ids: return empty + from litellm.proxy.proxy_server import prisma_client as _prisma_client + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache + + if _user_api_key_cache is None: + return empty + + models: List[str] = [] + mcp_ids: List[str] = [] + agent_ids: List[str] = [] + + for ag_id in access_group_ids: + try: + ag = await get_access_object( + access_group_id=ag_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + models.extend(getattr(ag, "access_model_names", [])) + mcp_ids.extend(getattr(ag, "access_mcp_server_ids", [])) + agent_ids.extend(getattr(ag, "access_agent_ids", [])) + except Exception: + verbose_proxy_logger.debug( + "Could not fetch access group %s for resource resolution", + ag_id, + ) + return { - "access_group_models": await _get_models_from_access_groups( - access_group_ids=access_group_ids, - ), - "access_group_mcp_server_ids": await _get_mcp_server_ids_from_access_groups( - access_group_ids=access_group_ids, - ), - "access_group_agent_ids": await _get_agent_ids_from_access_groups( - access_group_ids=access_group_ids, - ), + "access_group_models": list(set(models)), + "access_group_mcp_server_ids": list(set(mcp_ids)), + "access_group_agent_ids": list(set(agent_ids)), } @@ -3601,11 +3624,20 @@ async def list_team_v2( # Resolve resources inherited from access groups for each team if not use_deleted_table: - for team_item in team_list: - if isinstance(team_item, TeamListItem): - resolved = await _resolve_access_group_resources( - access_group_ids=team_item.access_group_ids, - ) + team_items_with_ag = [ + t for t in team_list + if isinstance(t, TeamListItem) and t.access_group_ids + ] + if team_items_with_ag: + results = await asyncio.gather( + *[ + _resolve_access_group_resources( + access_group_ids=t.access_group_ids, + ) + for t in team_items_with_ag + ] + ) + for team_item, resolved in zip(team_items_with_ag, results): team_item.access_group_models = resolved["access_group_models"] team_item.access_group_mcp_server_ids = resolved[ "access_group_mcp_server_ids" 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 366f659bda..a4b5e67760 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6491,3 +6491,166 @@ async def test_create_team_member_budget_table_with_duration(): assert budget_request.budget_duration == "30d" assert budget_request.max_budget == 20.0 assert result["metadata"]["team_member_budget_id"] == "budget-abc" + + +# --------------------------------------------------------------------------- +# Tests for _resolve_access_group_resources +# --------------------------------------------------------------------------- + + +class TestResolveAccessGroupResources: + """Tests for the single-pass access group resource resolution helper.""" + + @pytest.mark.asyncio + async def test_returns_empty_when_no_access_group_ids(self): + """None or empty list should return empty lists for all resource types.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + result_none = await _resolve_access_group_resources(access_group_ids=None) + assert result_none == { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + } + + result_empty = await _resolve_access_group_resources(access_group_ids=[]) + assert result_empty == { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + } + + @pytest.mark.asyncio + async def test_single_access_group(self): + """Single access group should return its resources.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + fake_ag = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="test-group", + access_model_names=["gpt-4", "claude-3"], + access_mcp_server_ids=["mcp-1"], + access_agent_ids=["agent-1", "agent-2"], + ) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1"], + ) + + assert sorted(result["access_group_models"]) == ["claude-3", "gpt-4"] + assert result["access_group_mcp_server_ids"] == ["mcp-1"] + assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + + @pytest.mark.asyncio + async def test_multiple_access_groups_deduplicates(self): + """Multiple access groups with overlapping resources should deduplicate.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + ag1 = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="group-1", + access_model_names=["gpt-4", "claude-3"], + access_mcp_server_ids=["mcp-1"], + access_agent_ids=["agent-1"], + ) + ag2 = LiteLLM_AccessGroupTable( + access_group_id="ag-2", + access_group_name="group-2", + access_model_names=["gpt-4", "gemini"], + access_mcp_server_ids=["mcp-1", "mcp-2"], + access_agent_ids=["agent-2"], + ) + + async def fake_get_access_object(access_group_id, **kwargs): + return {"ag-1": ag1, "ag-2": ag2}[access_group_id] + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + side_effect=fake_get_access_object, + ): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1", "ag-2"], + ) + + assert sorted(result["access_group_models"]) == ["claude-3", "gemini", "gpt-4"] + assert sorted(result["access_group_mcp_server_ids"]) == ["mcp-1", "mcp-2"] + assert sorted(result["access_group_agent_ids"]) == ["agent-1", "agent-2"] + + @pytest.mark.asyncio + async def test_missing_access_group_skipped(self): + """If an access group doesn't exist, it should be skipped gracefully.""" + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + ag1 = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="group-1", + access_model_names=["gpt-4"], + access_mcp_server_ids=[], + access_agent_ids=[], + ) + + async def fake_get_access_object(access_group_id, **kwargs): + if access_group_id == "ag-1": + return ag1 + raise HTTPException(status_code=404, detail="Not found") + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_access_object", + side_effect=fake_get_access_object, + ): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1", "ag-missing"], + ) + + assert result["access_group_models"] == ["gpt-4"] + assert result["access_group_mcp_server_ids"] == [] + assert result["access_group_agent_ids"] == [] + + @pytest.mark.asyncio + async def test_returns_empty_when_cache_unavailable(self): + """If user_api_key_cache is None, should return empty results.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_access_group_resources, + ) + + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", + None, + ): + result = await _resolve_access_group_resources( + access_group_ids=["ag-1"], + ) + + assert result == { + "access_group_models": [], + "access_group_mcp_server_ids": [], + "access_group_agent_ids": [], + }