[coding agent] Respect the --managed-identity-name flag on the command line

cli/azd/extensions/azure.coding-agent/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Release History
 
+## 0.5.2
+
+### Bugs fixed
+
+- Managed identity creation code was running multiple times, instead of a single time. This should not have broken anything, as the call itself is idempotent but calling it repeatedly is unnecessary.
+- `--managed-identity-name`, when specified, now assumes you're creating a credential, eliminating a question from the list.
+
 ## 0.5.1
 
 ### Bugs fixed

cli/azd/extensions/azure.coding-agent/extension.yaml

@@ -7,4 +7,4 @@ language: go
 namespace: coding-agent
 usage: azd coding-agent <command> [options]
 # NOTE: Make sure version.txt is in sync with this version.
-version: 0.5.1
+version: 0.5.2

cli/azd/extensions/azure.coding-agent/internal/cmd/coding_agent_config.go

@@ -50,6 +50,7 @@ var prBodyMD string
 
 const copilotEnv = "copilot"
 const readmeURL = "https://github.com/Azure/azure-dev/blob/main/cli/azd/extensions/azure.coding-agent/README.md"
+const defaultManagedIdentityName = "mi-copilot-coding-agent"
 
 type flagValues struct {
 	Debug               bool
@@ -89,7 +90,7 @@ func setupFlags(commandFlags *pflag.FlagSet) *flagValues {
 	commandFlags.StringVar(
 		&flagValues.ManagedIdentityName,
 		"managed-identity-name",
-		"mi-copilot-coding-agent",
+		defaultManagedIdentityName,
 		"The name to use for the managed identity, if created.",
 	)
 
@@ -523,30 +524,39 @@ func pickOrCreateMSI(ctx context.Context,
 
 	// ************************** Pick or create a new MSI **************************
 
-	// Prompt for pick or create a new MSI
-	selectedOption, err := prompter.Select(ctx, &azdext.SelectRequest{
-		Options: &azdext.SelectOptions{
-			Message: "Do you want to create a new Azure user-assigned managed identity or use an existing one?",
-			Choices: []*azdext.SelectChoice{
-				{Label: "Create new user-assigned managed identity"},
-				{Label: "Use existing user-assigned managed identity"},
+	shouldCreateIdentity := true
+
+	if identityName == defaultManagedIdentityName {
+		// Prompt for pick or create a new MSI
+		selectedOption, err := prompter.Select(ctx, &azdext.SelectRequest{
+			Options: &azdext.SelectOptions{
+				Message: "Do you want to create a new Azure user-assigned managed identity or use an existing one?",
+				Choices: []*azdext.SelectChoice{
+					{Label: "Create new user-assigned managed identity"},
+					{Label: "Use existing user-assigned managed identity"},
+				},
+				HelpMessage: "",
 			},
-			HelpMessage: "",
-		},
-	})
-	if err != nil {
-		//nolint:lll
-		return nil, fmt.Errorf(
-			"failed when prompting for managed identity option. Try logging in manually with 'azd auth login' before running this command. Error: %w",
-			err,
-		)
+		})
+
+		if err != nil {
+			//nolint:lll
+			return nil, fmt.Errorf(
+				"failed when prompting for managed identity option. Try logging in manually with 'azd auth login' before running this command. Error: %w",
+				err,
+			)
+		}
+
+		shouldCreateIdentity = *selectedOption.Value == 0
 	}
 
 	taskList := ux.NewTaskList(nil)
 
+	// this gets assigned when the task list runs - it's either a brand new identity, or an
+	// existing identity.
 	var managedIdentity rm_armmsi.Identity
 
-	if *selectedOption.Value == 0 {
+	if shouldCreateIdentity {
 		// pick a resource group and location for the new MSI
 		location, err := prompter.PromptLocation(ctx, &azdext.PromptLocationRequest{
 			AzureContext: &azdext.AzureContext{
@@ -646,21 +656,23 @@ func pickOrCreateMSI(ctx context.Context,
 		managedIdentity = msIdentities[*selectedOption.Value]
 	}
 
-	if err := taskList.Run(); err != nil {
-		return nil, err
-	}
-
 	roleNameStrings := strings.Join(roleNames, ", ")
-	parsedID, err := arm.ParseResourceID(*managedIdentity.ID)
-
-	if err != nil {
-		return nil, fmt.Errorf("invalid format for managed identity resource id: %w", err)
-	}
+	var resourceGroup string // filled out when we ensure role assignments
 
 	taskList.AddTask(ux.TaskOptions{
 		Title: fmt.Sprintf("Assigning roles (%s) to User Managed Identity (MSI)", roleNameStrings),
 		Action: func(spf ux.SetProgressFunc) (ux.TaskState, error) {
-			err := entraIDService.EnsureRoleAssignments(
+			// managedIdentity is filled out by an earlier step in this task list
+			// or chosen by the user when they say "use existing"
+			parsedID, err := arm.ParseResourceID(*managedIdentity.ID)
+
+			if err != nil {
+				return ux.Error, fmt.Errorf("invalid format for managed identity resource id: %w", err)
+			}
+
+			resourceGroup = parsedID.ResourceGroupName
+
+			err = entraIDService.EnsureRoleAssignments(
 				ctx,
 				subscriptionId,
 				roleNames,
@@ -687,7 +699,7 @@ func pickOrCreateMSI(ctx context.Context,
 
 	return &authConfiguration{
 		Name:           *managedIdentity.Name,
-		ResourceGroup:  parsedID.ResourceGroupName,
+		ResourceGroup:  resourceGroup,
 		TenantId:       *managedIdentity.Properties.TenantID,
 		SubscriptionId: subscriptionId,
 		ResourceID:     *managedIdentity.ID,

cli/azd/extensions/azure.coding-agent/internal/cmd/coding_agent_config_test.go

@@ -5,12 +5,15 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	azure_armmsi "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/msi/armmsi"
+	rm_armmsi "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/msi/armmsi"
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	exec "github.com/azure/azure-dev/cli/azd/pkg/exec"
+	"github.com/azure/azure-dev/cli/azd/pkg/graphsdk"
 	"github.com/azure/azure-dev/cli/azd/pkg/input"
 	azd_github "github.com/azure/azure-dev/cli/azd/pkg/tools/github"
 	"github.com/stretchr/testify/require"
@@ -21,6 +24,8 @@ import (
 
 //go:generate go tool mockgen -destination mocks_azdext_test.go -copyright_file ./testdata/mock_copyright.txt -package cmd github.com/azure/azure-dev/cli/azd/pkg/azdext PromptServiceClient
 
+//go:generate go tool mockgen -destination mocks_azd_entraid.go -copyright_file ./testdata/mock_copyright.txt -package cmd github.com/azure/azure-dev/cli/azd/pkg/entraid EntraIdService
+
 //go:generate go tool mockgen -destination mocks_azdexec_test.go -copyright_file ./testdata/mock_copyright.txt -package cmd github.com/azure/azure-dev/cli/azd/pkg/exec CommandRunner
 
 //go:generate go tool mockgen -destination mocks_azdinput_test.go -copyright_file ./testdata/mock_copyright.txt -package cmd github.com/azure/azure-dev/cli/azd/pkg/input Console
@@ -171,6 +176,185 @@ func TestCodingAgent_loginToGitHubIfNeeded(t *testing.T) {
 	})
 }
 
+func TestCodingAgent_pickOrCreateMSI(t *testing.T) {
+	const fakeSubscriptionID = "fake-subscription-id"
+	const fakeResourceGroup = "fake-resource-group"
+	const fakeLocation = "fake-location"
+	const fakeTenantID = "tenant-id"
+	const fakeClientID = "client-id"
+
+	setup := func(t *testing.T) (*MockPromptServiceClient, *MockazdMSIService, *MockEntraIdService, *MockresourceService) {
+		ctrl := gomock.NewController(t)
+		promptService := NewMockPromptServiceClient(ctrl)
+		msiService := NewMockazdMSIService(ctrl)
+		entraService := NewMockEntraIdService(ctrl)
+		resourceService := NewMockresourceService(ctrl)
+
+		promptService.EXPECT().
+			PromptLocation(gomock.Any(), gomock.Any()).
+			Return(&azdext.PromptLocationResponse{
+				Location: &azdext.Location{
+					Name: fakeLocation,
+				},
+			}, nil).
+			AnyTimes()
+
+		promptService.EXPECT().
+			PromptResourceGroup(gomock.Any(), gomock.Any()).
+			Return(&azdext.PromptResourceGroupResponse{
+				ResourceGroup: &azdext.ResourceGroup{
+					Name: fakeResourceGroup,
+				},
+			}, nil).
+			AnyTimes()
+
+		return promptService, msiService, entraService, resourceService
+	}
+
+	t.Run("UsingDefaultIdentity_createNewIdentity", func(t *testing.T) {
+		// the simplest "create" flow - they're going to create a new identity, using all the defaults.
+		promptService, msiService, entraService, resourceService := setup(t)
+
+		promptService.EXPECT().Select(gomock.Any(),
+			SelectMatcher{T: t, ExpectedQuestion: "Do you want to create a new Azure user-assigned managed identity or use an existing one?"}).
+			Return(&azdext.SelectResponse{
+				Value: to.Ptr(int32(0)), // ie, "create a new one"
+			}, nil)
+
+		msiService.EXPECT().
+			CreateUserIdentity(gomock.Any(), fakeSubscriptionID, fakeResourceGroup, fakeLocation, defaultManagedIdentityName).
+			Return(azure_armmsi.Identity{
+				Properties: &azure_armmsi.UserAssignedIdentityProperties{
+					PrincipalID: to.Ptr("principal-id"),
+					ClientID:    to.Ptr(fakeClientID),
+					TenantID:    to.Ptr(fakeTenantID),
+				},
+				Name: to.Ptr(string(defaultManagedIdentityName)),
+				ID:   to.Ptr(fmt.Sprintf("/subscriptions/%s/resourcegroups/%s/providers/Microsoft.ManagedIdentity/userAssignedIdentities/%s", fakeSubscriptionID, fakeResourceGroup, defaultManagedIdentityName)),
+			}, nil)
+
+		entraService.EXPECT().EnsureRoleAssignments(gomock.Any(),
+			fakeSubscriptionID,
+			[]string{"Reader"},
+			&graphsdk.ServicePrincipal{
+				Id:          to.Ptr("principal-id"),
+				DisplayName: defaultManagedIdentityName,
+			}, gomock.Any()).
+			Return(nil)
+
+		authConfig, err := pickOrCreateMSI(context.Background(),
+			promptService,
+			msiService,
+			entraService,
+			resourceService,
+			defaultManagedIdentityName,
+			fakeSubscriptionID,
+			[]string{"Reader"})
+		require.NoError(t, err)
+
+		require.Equal(t, fakeClientID, authConfig.ClientId)
+		require.Equal(t, defaultManagedIdentityName, authConfig.Name)
+		require.Equal(t, fakeTenantID, authConfig.TenantId)
+		require.Equal(t, fakeResourceGroup, authConfig.ResourceGroup)
+	})
+
+	t.Run("UsingCustomManagedIdentityNameAndRoles_assumesCreate", func(t *testing.T) {
+		// similar to the create flow above, except that specifying a custom name for the identity
+		// automaticaly chooses the "create" path.
+		promptService, msiService, entraService, resourceService := setup(t)
+		const customManagedIdentityName = "identity-name"
+
+		msiService.EXPECT().
+			CreateUserIdentity(gomock.Any(), fakeSubscriptionID, fakeResourceGroup, fakeLocation, customManagedIdentityName).
+			Return(azure_armmsi.Identity{
+				Properties: &azure_armmsi.UserAssignedIdentityProperties{
+					PrincipalID: to.Ptr("principal-id"),
+					ClientID:    to.Ptr(fakeClientID),
+					TenantID:    to.Ptr(fakeTenantID),
+				},
+				Name: to.Ptr(customManagedIdentityName),
+				ID:   to.Ptr(fmt.Sprintf("/subscriptions/%s/resourcegroups/%s/providers/Microsoft.ManagedIdentity/userAssignedIdentities/%s", fakeSubscriptionID, fakeResourceGroup, customManagedIdentityName)),
+			}, nil)
+
+		entraService.EXPECT().EnsureRoleAssignments(gomock.Any(), fakeSubscriptionID, []string{"custom-role-name"}, &graphsdk.ServicePrincipal{
+			Id:          to.Ptr("principal-id"),
+			DisplayName: customManagedIdentityName,
+		}, gomock.Any()).
+			Return(nil)
+
+		authConfig, err := pickOrCreateMSI(context.Background(),
+			promptService,
+			msiService,
+			entraService,
+			resourceService,
+			customManagedIdentityName,
+			fakeSubscriptionID,
+			[]string{"custom-role-name"})
+		require.NoError(t, err)
+
+		require.Equal(t, fakeClientID, authConfig.ClientId)
+		require.Equal(t, customManagedIdentityName, authConfig.Name)
+		require.Equal(t, fakeTenantID, authConfig.TenantId)
+		require.Equal(t, fakeResourceGroup, authConfig.ResourceGroup)
+	})
+
+	t.Run("UsingDefaultIdentity_useExisting", func(t *testing.T) {
+		// uses an existing identity - no identity should be created!
+		promptService, msiService, entraService, resourceService := setup(t)
+
+		promptService.EXPECT().Select(gomock.Any(),
+			SelectMatcher{T: t, ExpectedQuestion: "Do you want to create a new Azure user-assigned managed identity or use an existing one?"}).
+			Return(&azdext.SelectResponse{
+				Value: to.Ptr(int32(1)), // ie, "use an existing one"
+			}, nil)
+
+		msiService.EXPECT().
+			ListUserIdentities(gomock.Any(), fakeSubscriptionID).
+			Return([]rm_armmsi.Identity{
+				{
+					Properties: &azure_armmsi.UserAssignedIdentityProperties{
+						PrincipalID: to.Ptr("principal-id"),
+						ClientID:    to.Ptr(fakeClientID),
+						TenantID:    to.Ptr(fakeTenantID),
+					},
+					Name:     to.Ptr(string(defaultManagedIdentityName)),
+					ID:       to.Ptr(fmt.Sprintf("/subscriptions/%s/resourcegroups/%s/providers/Microsoft.ManagedIdentity/userAssignedIdentities/%s", fakeSubscriptionID, fakeResourceGroup, defaultManagedIdentityName)),
+					Location: to.Ptr(fakeLocation),
+				},
+			}, nil)
+
+		promptService.EXPECT().
+			Select(gomock.Any(), SelectMatcher{T: t, ExpectedQuestion: "Select an existing User Managed Identity (MSI) to use"}).
+			Return(&azdext.SelectResponse{
+				Value: to.Ptr(int32(0)), // pick the first ID in the list (also, the only ID I returned)
+			}, nil)
+
+		entraService.EXPECT().EnsureRoleAssignments(gomock.Any(),
+			fakeSubscriptionID,
+			[]string{"Reader"},
+			&graphsdk.ServicePrincipal{
+				Id:          to.Ptr("principal-id"),
+				DisplayName: defaultManagedIdentityName,
+			}, gomock.Any()).
+			Return(nil)
+
+		authConfig, err := pickOrCreateMSI(context.Background(),
+			promptService,
+			msiService,
+			entraService,
+			resourceService,
+			defaultManagedIdentityName,
+			fakeSubscriptionID,
+			[]string{"Reader"})
+		require.NoError(t, err)
+
+		require.Equal(t, fakeClientID, authConfig.ClientId)
+		require.Equal(t, defaultManagedIdentityName, authConfig.Name)
+		require.Equal(t, fakeTenantID, authConfig.TenantId)
+		require.Equal(t, fakeResourceGroup, authConfig.ResourceGroup)
+	})
+}
+
 type FedCredentialMatcher struct {
 	T        *testing.T
 	Expected azure_armmsi.FederatedIdentityCredential
@@ -197,3 +381,18 @@ func (m GitHubEnvMatcher) Matches(x any) bool {
 
 // String describes what the matcher matches.
 func (m GitHubEnvMatcher) String() string { return "Checks copilot env was specified" }
+
+type SelectMatcher struct {
+	T                *testing.T
+	ExpectedQuestion string
+}
+
+func (m SelectMatcher) Matches(x any) bool {
+	req := x.(*azdext.SelectRequest)
+	require.Equal(m.T, m.ExpectedQuestion, req.Options.Message)
+	return true
+}
+
+func (m SelectMatcher) String() string {
+	return "Checks that the select prompt was the one we're expecting"
+}