Miscellaneous improvements related to durable task hub

package.json

@@ -1483,9 +1483,9 @@
         "@azure/storage-blob": "^12.5.0",
         "@microsoft/vscode-azext-azureappservice": "^3.6.4",
         "@microsoft/vscode-azext-azureappsettings": "^0.2.8",
-        "@microsoft/vscode-azext-azureutils": "^3.4.7",
+        "@microsoft/vscode-azext-azureutils": "^3.4.8",
         "@microsoft/vscode-azext-utils": "^3.3.3",
-        "@microsoft/vscode-azureresources-api": "^2.0.4",
+        "@microsoft/vscode-azureresources-api": "^2.5.1",
         "@microsoft/vscode-container-client": "^0.1.2",
         "cross-fetch": "^4.0.0",
         "escape-string-regexp": "^4.0.0",

src/commands/appSettings/connectionSettings/durableTaskScheduler/DTSConnectionSetSettingStep.ts

@@ -9,7 +9,6 @@ import { CodeAction, ConnectionKey, hostFileName } from '../../../../constants';
 import { ext } from '../../../../extensionVariables';
 import { type IDTSTaskJson, type IHostJsonV2 } from '../../../../funcConfig/host';
 import { localize } from '../../../../localize';
-import { clientIdKey } from '../../../durableTaskScheduler/copySchedulerConnectionString';
 import { notifyFailedToConfigureHost } from '../notifyFailedToConfigureHost';
 import { setLocalSetting } from '../setConnectionSetting';
 import { type IDTSAzureConnectionWizardContext, type IDTSConnectionWizardContext } from './IDTSConnectionWizardContext';
@@ -36,22 +35,15 @@ export class DTSConnectionSetSettingStep<T extends IDTSConnectionWizardContext |
             context.newDTSConnectionSettingKey = ConnectionKey.DTS;
         }
 
-        const newDTSConnectionSettingKey = nonNullProp(context, 'newDTSConnectionSettingKey');
-        let newDTSConnectionSettingValue = nonNullProp(context, 'newDTSConnectionSettingValue');
-
-        // Todo: Move this to `DurableTaskSchedulerGetConnectionStep` when we upgrade the azure package for new identity logic
-        if ((context as IDTSAzureConnectionWizardContext).managedIdentity) {
-            newDTSConnectionSettingValue = newDTSConnectionSettingValue.replace(clientIdKey, (context as IDTSAzureConnectionWizardContext).managedIdentity?.clientId ?? clientIdKey);
-        }
-
         if (context.action === CodeAction.Debug) {
-            await setLocalSetting(context, newDTSConnectionSettingKey, newDTSConnectionSettingValue);
+            await setLocalSetting(
+                context,
+                nonNullProp(context, 'newDTSConnectionSettingKey'),
+                nonNullProp(context, 'newDTSConnectionSettingValue'),
+            );
         } else {
             // No further action required
         }
-
-        context.newDTSConnectionSettingValue = newDTSConnectionSettingValue;
-        context.valuesToMask.push(context.newDTSConnectionSettingValue);
     }
 
     public shouldExecute(context: T): boolean {

src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DTSStartingResourcesLogStep.ts

@@ -12,16 +12,12 @@ const startingResourcesContext: string = 'startingResourcesLogStepItem';
 
 export class DTSStartingResourcesLogStep<T extends IDTSAzureConnectionWizardContext> extends AzureWizardPromptStep<T> {
     public hideStepCount: boolean = true;
-    protected hasLogged: boolean = false;
 
     public async configureBeforePrompt(context: T): Promise<void> {
-        if (this.hasLogged) {
-            return;
-        }
-
         if (context.resourceGroup) {
             prependOrInsertAfterLastInfoChild(context,
                 new ActivityChildItem({
+                    stepId: this.id,
                     contextValue: createContextValue([startingResourcesContext, activityInfoContext]),
                     label: localize('useResourceGroup', 'Use resource group "{0}"', context.resourceGroup.name),
                     activityType: ActivityChildType.Info,
@@ -34,6 +30,7 @@ export class DTSStartingResourcesLogStep<T extends IDTSAzureConnectionWizardCont
         if (context.site) {
             prependOrInsertAfterLastInfoChild(context,
                 new ActivityChildItem({
+                    stepId: this.id,
                     label: localize('useFunctionApp', 'Use function app "{0}"', context.site.fullName),
                     contextValue: createContextValue([startingResourcesContext, activityInfoContext]),
                     activityType: ActivityChildType.Info,
@@ -46,6 +43,7 @@ export class DTSStartingResourcesLogStep<T extends IDTSAzureConnectionWizardCont
         if (context.dts) {
             prependOrInsertAfterLastInfoChild(context,
                 new ActivityChildItem({
+                    stepId: this.id,
                     label: localize('useDTS', 'Use durable task scheduler "{0}"', context.dts.name),
                     contextValue: createContextValue([startingResourcesContext, activityInfoContext]),
                     activityType: ActivityChildType.Info,
@@ -58,6 +56,7 @@ export class DTSStartingResourcesLogStep<T extends IDTSAzureConnectionWizardCont
         if (context.dtsHub) {
             prependOrInsertAfterLastInfoChild(context,
                 new ActivityChildItem({
+                    stepId: this.id,
                     label: localize('useDTSHub', 'Use durable task hub "{0}"', context.dtsHub.name),
                     contextValue: createContextValue([startingResourcesContext, activityInfoContext]),
                     activityType: ActivityChildType.Info,
@@ -74,8 +73,6 @@ export class DTSStartingResourcesLogStep<T extends IDTSAzureConnectionWizardCont
         if (context.newDTSHubConnectionSettingKey) {
             ext.outputChannel.appendLog(localize('dtsHubConnectionKey', 'Using DTS hub host connection key "{0}"', context.newDTSHubConnectionSettingKey));
         }
-
-        this.hasLogged = true;
     }
 
     public async prompt(): Promise<void> {

src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DurableTaskHubListStep.ts

@@ -3,11 +3,13 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { parseAzureResourceId } from '@microsoft/vscode-azext-azureutils';
-import { AzureWizardPromptStep, nonNullProp, type IAzureQuickPickItem, type IWizardOptions } from '@microsoft/vscode-azext-utils';
+import { CommonRoleDefinitions, createAuthorizationManagementClient, createRoleId, parseAzureResourceId, RoleAssignmentExecuteStep, uiUtils, type Role } from '@microsoft/vscode-azext-azureutils';
+import { ActivityChildItem, ActivityChildType, activitySuccessContext, activitySuccessIcon, AzureWizardPromptStep, createContextValue, nonNullProp, type AzureWizardExecuteStep, type IAzureQuickPickItem, type IWizardOptions } from '@microsoft/vscode-azext-utils';
 import { localSettingsDescription } from '../../../../../constants-nls';
+import { ext } from '../../../../../extensionVariables';
 import { localize } from '../../../../../localize';
 import { HttpDurableTaskSchedulerClient, type DurableTaskHubResource, type DurableTaskSchedulerClient } from '../../../../../tree/durableTaskScheduler/DurableTaskSchedulerClient';
+import { FunctionAppUserAssignedIdentitiesListStep } from '../../../../identity/FunctionAppUserAssignedIdentitiesListStep';
 import { type IDTSAzureConnectionWizardContext } from '../IDTSConnectionWizardContext';
 import { DurableTaskHubCreateStep } from './DurableTaskHubCreateStep';
 import { DurableTaskHubNameStep } from './DurableTaskHubNameStep';
@@ -22,7 +24,7 @@ export class DurableTaskHubListStep<T extends IDTSAzureConnectionWizardContext>
 
     public async prompt(context: T): Promise<void> {
         context.dtsHub = (await context.ui.showQuickPick(await this.getPicks(context), {
-            placeHolder: localize('selectTaskScheduler', 'Select a Durable Task Scheduler'),
+            placeHolder: localize('selectTaskScheduler', 'Select a durable task hub'),
         })).data;
 
         if (context.dtsHub) {
@@ -35,17 +37,6 @@ export class DurableTaskHubListStep<T extends IDTSAzureConnectionWizardContext>
         return !context.dtsHub;
     }
 
-    public async getSubWizard(context: T): Promise<IWizardOptions<T> | undefined> {
-        if (context.dtsHub) {
-            return undefined;
-        }
-
-        return {
-            promptSteps: [new DurableTaskHubNameStep(this.schedulerClient)],
-            executeSteps: [new DurableTaskHubCreateStep(this.schedulerClient)],
-        };
-    }
-
     private async getPicks(context: T): Promise<IAzureQuickPickItem<DurableTaskHubResource | undefined>[]> {
         const taskHubs: DurableTaskHubResource[] = context.dts ?
             await this.schedulerClient.getSchedulerTaskHubs(nonNullProp(context, 'subscription'), parseAzureResourceId(context.dts.id).resourceGroup, context.dts.name) : [];
@@ -66,4 +57,80 @@ export class DurableTaskHubListStep<T extends IDTSAzureConnectionWizardContext>
             }),
         ];
     }
+
+    public async getSubWizard(context: T): Promise<IWizardOptions<T> | undefined> {
+        const promptSteps: AzureWizardPromptStep<T>[] = [];
+        const executeSteps: AzureWizardExecuteStep<T>[] = [];
+
+        if (!context.dtsHub) {
+            promptSteps.push(new DurableTaskHubNameStep(this.schedulerClient));
+            executeSteps.push(new DurableTaskHubCreateStep(this.schedulerClient));
+        }
+
+        const dtsContributorRole: Role = {
+            scopeId: context.dtsHub?.id,
+            roleDefinitionId: createRoleId(context.subscriptionId, CommonRoleDefinitions.durableTaskDataContributor),
+            roleDefinitionName: CommonRoleDefinitions.durableTaskDataContributor.roleName,
+        };
+
+        promptSteps.push(new FunctionAppUserAssignedIdentitiesListStep(dtsContributorRole /** targetRole */, { identityAssignStepPriority: 180 }));
+        executeSteps.push(new RoleAssignmentExecuteStep(this.getDTSRoleAssignmentCallback(context, dtsContributorRole), { priority: 190 }));
+
+        return { promptSteps, executeSteps };
+    }
+
+    private getDTSRoleAssignmentCallback(context: T, role: Role): () => Promise<Role[]> {
+        return async () => {
+            const roleAssignment: Role = {
+                ...role,
+                // This id may be missing when the role is initially passed in,
+                // but by the time we run the step, we should have the populated id ready.
+                scopeId: context.dtsHub?.id,
+            };
+
+            if (!roleAssignment.scopeId) {
+                return [];
+            }
+
+            const amClient = await createAuthorizationManagementClient(context);
+
+            let hasRoleAssignment: boolean = false;
+            if (context.dtsHub) {
+                const taskHubRoleAssignments = await uiUtils.listAllIterator(amClient.roleAssignments.listForScope(
+                    context.dtsHub.id,
+                    {
+                        // $filter=principalId eq {id}
+                        filter: `principalId eq '{${context.managedIdentity?.principalId}}'`,
+                    }
+                ));
+                hasRoleAssignment = taskHubRoleAssignments.some(r => !!r.roleDefinitionId?.endsWith(role.roleDefinitionId));
+            }
+
+            if (!hasRoleAssignment && context.dts) {
+                const taskSchedulerRoleAssignments = await uiUtils.listAllIterator(amClient.roleAssignments.listForScope(
+                    context.dts.id,
+                    {
+                        // $filter=principalId eq {id}
+                        filter: `principalId eq '{${context.managedIdentity?.principalId}}'`,
+                    }
+                ));
+                hasRoleAssignment = taskSchedulerRoleAssignments.some(r => !!r.roleDefinitionId?.endsWith(role.roleDefinitionId));
+            }
+
+            if (hasRoleAssignment) {
+                context.activityChildren?.push(
+                    new ActivityChildItem({
+                        label: localize('verifyIdentityWithRoleLabel', 'Verify identity "{0}" has role "{1}"', context.managedIdentity?.name, role.roleDefinitionName),
+                        description: '0s',
+                        contextValue: createContextValue(['roleAssignmentExecuteStepItem', activitySuccessContext]),
+                        activityType: ActivityChildType.Success,
+                        iconPath: activitySuccessIcon,
+                    }),
+                );
+                ext.outputChannel.appendLog(localize('verifyIdentity', 'Successfully verified identity "{0}" has role "{1}".', context.managedIdentity?.name, role.roleDefinitionName));
+            }
+
+            return hasRoleAssignment ? [] : [roleAssignment];
+        };
+    }
 }

src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DurableTaskHubNameStep.ts

@@ -20,7 +20,7 @@ export class DurableTaskHubNameStep<T extends IDTSAzureConnectionWizardContext>
     public async prompt(context: T): Promise<void> {
         context.newDTSHubName = (await context.ui.showInputBox({
             prompt: localize('taskSchedulerName', 'Enter a name for the durable task hub'),
-            value: context.suggestedDTSHubNameLocalSettings,
+            value: context.suggestedDTSHubNameLocalSettings ?? 'default',
             validateInput: this.validateInput,
             asyncValidationTask: (name: string) => this.validateNameAvailable(context, name),
         })).trim();

src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DurableTaskSchedulerGetConnectionStep.ts

@@ -4,14 +4,19 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { AzureWizardExecuteStep } from '@microsoft/vscode-azext-utils';
-import { getSchedulerConnectionString, SchedulerAuthenticationType } from '../../../../durableTaskScheduler/copySchedulerConnectionString';
+import { clientIdKey, getSchedulerConnectionString, SchedulerAuthenticationType } from '../../../../durableTaskScheduler/copySchedulerConnectionString';
 import { type IDTSAzureConnectionWizardContext } from '../IDTSConnectionWizardContext';
 
 export class DurableTaskSchedulerGetConnectionStep<T extends IDTSAzureConnectionWizardContext> extends AzureWizardExecuteStep<T> {
     public priority: number = 200;
 
     public async execute(context: T): Promise<void> {
         context.newDTSConnectionSettingValue = getSchedulerConnectionString(context.dts?.properties.endpoint ?? '', SchedulerAuthenticationType.UserAssignedIdentity);
+
+        if (context.managedIdentity) {
+            context.newDTSConnectionSettingValue = context.newDTSConnectionSettingValue.replace(clientIdKey, context.managedIdentity?.clientId ?? clientIdKey);
+        }
+
         context.valuesToMask.push(context.newDTSConnectionSettingValue);
     }
 

src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DurableTaskSchedulerListStep.ts

@@ -3,12 +3,11 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { CommonRoleDefinitions, createAuthorizationManagementClient, createRoleId, LocationListStep, parseAzureResourceId, RoleAssignmentExecuteStep, uiUtils, type ILocationWizardContext, type Role } from '@microsoft/vscode-azext-azureutils';
+import { LocationListStep, parseAzureResourceId, type ILocationWizardContext } from '@microsoft/vscode-azext-azureutils';
 import { AzureWizardPromptStep, nonNullProp, type AzureWizardExecuteStep, type IAzureQuickPickItem, type IWizardOptions } from '@microsoft/vscode-azext-utils';
 import { localSettingsDescription } from '../../../../../constants-nls';
 import { localize } from '../../../../../localize';
 import { HttpDurableTaskSchedulerClient, type DurableTaskSchedulerClient, type DurableTaskSchedulerResource } from '../../../../../tree/durableTaskScheduler/DurableTaskSchedulerClient';
-import { FunctionAppUserAssignedIdentitiesListStep } from '../../../../identity/FunctionAppUserAssignedIdentitiesListStep';
 import { type IDTSAzureConnectionWizardContext } from '../IDTSConnectionWizardContext';
 import { DurableTaskHubListStep } from './DurableTaskHubListStep';
 import { DurableTaskSchedulerCreateStep } from './DurableTaskSchedulerCreateStep';
@@ -52,43 +51,7 @@ export class DurableTaskSchedulerListStep<T extends IDTSAzureConnectionWizardCon
             promptSteps.push(new DurableTaskHubListStep(this._schedulerClient));
         }
 
-        const dtsContributorRole: Role = {
-            scopeId: context.dts?.id,
-            roleDefinitionId: createRoleId(context.subscriptionId, CommonRoleDefinitions.durableTaskDataContributor),
-            roleDefinitionName: CommonRoleDefinitions.durableTaskDataContributor.roleName,
-        };
-
-        promptSteps.push(new FunctionAppUserAssignedIdentitiesListStep(dtsContributorRole /** targetRole */, { identityAssignStepPriority: 180 }));
-        executeSteps.push(new RoleAssignmentExecuteStep(getDTSRoleAssignmentCallback(context, dtsContributorRole), { priority: 190 }));
-
         return { promptSteps, executeSteps };
-
-        function getDTSRoleAssignmentCallback(context: T, role: Role): () => Promise<Role[]> {
-            return async () => {
-                const roleAssignment: Role = {
-                    ...role,
-                    // This id may be missing when the role is initially passed in,
-                    // but by the time we run the step, we should have the populated id ready.
-                    scopeId: context.dts?.id,
-                };
-
-                if (!roleAssignment.scopeId) {
-                    return [];
-                }
-
-                const amClient = await createAuthorizationManagementClient(context);
-                const roleAssignments = await uiUtils.listAllIterator(amClient.roleAssignments.listForScope(
-                    roleAssignment.scopeId,
-                    {
-                        // $filter=principalId eq {id}
-                        filter: `principalId eq '{${context.managedIdentity?.principalId}}'`,
-                    }
-                ));
-
-                const hasRoleAssignment = roleAssignments.some(r => !!r.roleDefinitionId?.endsWith(role.roleDefinitionId));
-                return hasRoleAssignment ? [] : [roleAssignment];
-            };
-        }
     }
 
     private async getPicks(context: T): Promise<IAzureQuickPickItem<DurableTaskSchedulerResource | undefined>[]> {