Conslidate DTS logic

Author: MicroFish91

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';
@@ -37,12 +36,7 @@ export class DTSConnectionSetSettingStep<T extends IDTSConnectionWizardContext |
         }
 
         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);
-        }
+        const newDTSConnectionSettingValue = nonNullProp(context, 'newDTSConnectionSettingValue');
 
         if (context.action === CodeAction.Debug) {
             await setLocalSetting(context, newDTSConnectionSettingKey, newDTSConnectionSettingValue);

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 as IDTSAzureConnectionWizardContext).managedIdentity?.clientId ?? clientIdKey);
+        }
+
         context.valuesToMask.push(context.newDTSConnectionSettingValue);
     }