utils: Fix 'Loading...' quick pick remaining visible during create callback (#2227)

* Fix 'Loading...' quick pick remaining visible during create callback

When a user selects '+ Create new...' in the tree item picker, the outer
wizard's loading quick pick re-appears after each inner wizard prompt
resolves (via onDidFinishPrompt) and stays visible during the execute
phase when no other quick pick overlays it.

Add a suppressLoadingPrompt context flag that is set during the create
callback execution in CompatibilityRecursiveQuickPickStep. AzureWizard
checks this flag before re-showing the loading quick pick and before
cancelling the wizard on hide.

Fixes microsoft/vscode-azureappservice#2833

* Address PR review feedback for suppressLoadingPrompt

- Save/restore previous suppressLoadingPrompt value instead of forcing false
- Fix misleading test comment to describe actual failure path
- Update doc comment to accurately reflect suppression scope

* Extract repeated IInternalActionContext cast into local variable

Author: alexweininger

utils/package-lock.json

@@ -1,7 +1,7 @@
 {
     "name": "@microsoft/vscode-azext-utils",
     "author": "Microsoft Corporation",
-    "version": "4.0.4",
+    "version": "4.0.5",
     "description": "Common UI tools for developing Azure extensions for VS Code",
     "tags": [
         "azure",

utils/package.json

@@ -9,6 +9,7 @@ import * as types from "../../../../index";
 import { NoResourceFoundError, UserCancelledError } from "../../../errors";
 import { AzExtTreeItem } from "../../../tree/AzExtTreeItem";
 import { isAzExtParentTreeItem, isAzExtTreeItem } from "../../../tree/isAzExtTreeItem";
+import { IInternalActionContext } from "../../../userInput/IInternalActionContext";
 import { getLastNode } from "../../getLastNode";
 import { CompatibilityContextValueQuickPickStep } from './CompatibilityContextValueQuickPickStep';
 
@@ -65,19 +66,31 @@ export class CompatibilityRecursiveQuickPickStep<TContext extends types.QuickPic
 
                 // context passed to callback must have the same `ui` as the wizardContext
                 // to prevent the wizard from being cancelled unexpectedly
-                const createdTreeItem = await callback?.(wizardContext) as AzExtTreeItem;
-
-                // convert created AzExtTreeItem to a BranchDataItem so the wizard can get its children in the next step
-                const picks = await this.getPicks(wizardContext) as types.IAzureQuickPickItem<unknown>[];
-                const createdPick = picks.find((pick) => {
-                    return (pick.data as Wrapper).unwrap<AzExtTreeItem>().fullId === createdTreeItem.fullId;
-                });
-
-                if (createdPick) {
-                    return createdPick.data;
+                //
+                // Suppress the outer wizard's loading quick pick during the create callback.
+                // Without this, the "Loading..." placeholder re-appears after each inner wizard
+                // prompt resolves (via onDidFinishPrompt) and stays visible during the execute
+                // phase when no other quick pick overlays it.
+                const internalContext = wizardContext as unknown as IInternalActionContext;
+                const previousSuppressLoadingPrompt = internalContext.suppressLoadingPrompt;
+                internalContext.suppressLoadingPrompt = true;
+                try {
+                    const createdTreeItem = await callback?.(wizardContext) as AzExtTreeItem;
+
+                    // convert created AzExtTreeItem to a BranchDataItem so the wizard can get its children in the next step
+                    const picks = await this.getPicks(wizardContext) as types.IAzureQuickPickItem<unknown>[];
+                    const createdPick = picks.find((pick) => {
+                        return (pick.data as Wrapper).unwrap<AzExtTreeItem>().fullId === createdTreeItem.fullId;
+                    });
+
+                    if (createdPick) {
+                        return createdPick.data;
+                    }
+
+                    throw new UserCancelledError();
+                } finally {
+                    internalContext.suppressLoadingPrompt = previousSuppressLoadingPrompt;
                 }
-
-                throw new UserCancelledError();
             }
 
             return selected.data;

utils/src/pickTreeItem/contextValue/compatibility/CompatibilityRecursiveQuickPickStep.ts

@@ -8,6 +8,14 @@ import * as types from '../../index';
 
 export interface IInternalActionContext extends types.IActionContext {
     ui: types.IAzureUserInput & { wizard?: IInternalAzureWizard, isPrompting?: boolean, isTesting?: boolean };
+
+    /**
+     * When true, the outer wizard's "Loading..." quick pick will not be
+     * re-shown after inner prompts resolve and cancel-on-hide will be
+     * suppressed. Used internally to prevent the loading placeholder from
+     * remaining visible during a nested create-child wizard execution.
+     */
+    suppressLoadingPrompt?: boolean;
 }
 
 export interface IInternalAzureWizard {

utils/src/userInput/IInternalActionContext.ts

@@ -119,7 +119,7 @@ export class AzureWizard<T extends (IInternalActionContext & Partial<types.Execu
                     if (loadingQuickPick) {
                         disposables.push(loadingQuickPick?.onDidHide(() => {
                             // Avoid issuing cancels during tests - see: https://github.com/microsoft/vscode-azuretools/pull/1807
-                            if (!this._context.ui.isPrompting && !this._context.ui.isTesting) {
+                            if (!this._context.ui.isPrompting && !this._context.ui.isTesting && !this._context.suppressLoadingPrompt) {
                                 this._cancellationTokenSource.cancel();
                             }
                         }));
@@ -128,7 +128,9 @@ export class AzureWizard<T extends (IInternalActionContext & Partial<types.Execu
                     disposables.push(this._context.ui.onDidFinishPrompt((result) => {
                         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
                         step!.prompted = true;
-                        loadingQuickPick?.show();
+                        if (!this._context.suppressLoadingPrompt) {
+                            loadingQuickPick?.show();
+                        }
                         if (typeof result.value === 'string' && !result.matchesDefault && this.currentStepId && !step?.supportsDuplicateSteps) {
                             this._cachedInputBoxValues[this.currentStepId] = result.value;
                         }

utils/src/wizard/AzureWizard.ts

@@ -663,4 +663,97 @@ suite("AzureWizard tests", () => {
             { quickPick1: 'Pick 1', quickPick2: 'Pick 2', inputBox1: 'testValue2', subQuickPick1: 'Pick 3', execute1: 'executeValue1' }
         );
     });
+
+    test("showLoadingPrompt does not break normal wizard flow", async () => {
+        const testUserInput: TestUserInput = new TestUserInput(vscode);
+        const context: ITestWizardContext = { telemetry: { properties: {}, measurements: {} }, errorHandling: { issueProperties: {} }, ui: testUserInput, valuesToMask: [] };
+
+        const wizard = new AzureWizard<ITestWizardContext>(context, {
+            promptSteps: [new QuickPickStep1(), new InputBoxStep1()],
+            executeSteps: [new ExecuteStep1()],
+            showLoadingPrompt: true,
+        });
+
+        await testUserInput.runWithInputs(['Pick 1', 'testValue'], async () => {
+            await wizard.prompt();
+        });
+        await wizard.execute();
+
+        assert.strictEqual(context.quickPick1, 'Pick 1');
+        assert.strictEqual(context.inputBox1, 'testValue');
+        assert.strictEqual(context.execute1, 'executeValue1');
+    });
+
+    test("suppressLoadingPrompt prevents loading prompt re-show without breaking wizard", async () => {
+        const testUserInput: TestUserInput = new TestUserInput(vscode);
+        const context: ITestWizardContext = { telemetry: { properties: {}, measurements: {} }, errorHandling: { issueProperties: {} }, ui: testUserInput, valuesToMask: [] };
+
+        // Simulate what CompatibilityRecursiveQuickPickStep does:
+        // set suppressLoadingPrompt before a prompt step that triggers an inner wizard
+        const stepThatSuppresses = new class extends AzureWizardPromptStep<ITestWizardContext> {
+            async prompt(wizardContext: ITestWizardContext): Promise<void> {
+                // Set the flag like CompatibilityRecursiveQuickPickStep does before the create callback
+                wizardContext.suppressLoadingPrompt = true;
+                try {
+                    // Simulate inner wizard prompts (these would fire onDidFinishPrompt)
+                    wizardContext['innerResult'] = (await wizardContext.ui.showQuickPick(
+                        [{ label: 'Inner Pick 1' }, { label: 'Inner Pick 2' }],
+                        { placeHolder: 'Inner wizard pick' }
+                    )).label;
+                    wizardContext['innerInput'] = await wizardContext.ui.showInputBox({ prompt: 'Inner wizard input' });
+                } finally {
+                    wizardContext.suppressLoadingPrompt = false;
+                }
+            }
+            shouldPrompt(): boolean { return true; }
+        };
+
+        const wizard = new AzureWizard<ITestWizardContext>(context, {
+            promptSteps: [stepThatSuppresses, new QuickPickStep1()],
+            executeSteps: [new ExecuteStep1()],
+            showLoadingPrompt: true,
+        });
+
+        await testUserInput.runWithInputs(['Inner Pick 1', 'innerTestValue', 'Pick 2'], async () => {
+            await wizard.prompt();
+        });
+        await wizard.execute();
+
+        assert.strictEqual(context.innerResult, 'Inner Pick 1');
+        assert.strictEqual(context.innerInput, 'innerTestValue');
+        assert.strictEqual(context.quickPick1, 'Pick 2');
+        assert.strictEqual(context.execute1, 'executeValue1');
+        // Flag should be reset after the step completes
+        assert.strictEqual(context.suppressLoadingPrompt, false);
+    });
+
+    test("suppressLoadingPrompt is reset even if inner prompts throw", async () => {
+        const testUserInput: TestUserInput = new TestUserInput(vscode);
+        const context: ITestWizardContext = { telemetry: { properties: {}, measurements: {} }, errorHandling: { issueProperties: {} }, ui: testUserInput, valuesToMask: [] };
+
+        const stepThatSuppressesAndFails = new class extends AzureWizardPromptStep<ITestWizardContext> {
+            async prompt(wizardContext: ITestWizardContext): Promise<void> {
+                wizardContext.suppressLoadingPrompt = true;
+                try {
+                    // Ensure suppressLoadingPrompt is reset even when the step throws
+                    throw new Error('simulated inner wizard failure');
+                } finally {
+                    wizardContext.suppressLoadingPrompt = false;
+                }
+            }
+            shouldPrompt(): boolean { return true; }
+        };
+
+        const wizard = new AzureWizard<ITestWizardContext>(context, {
+            promptSteps: [stepThatSuppressesAndFails],
+            showLoadingPrompt: true,
+        });
+
+        await assert.rejects(async () => {
+            await wizard.prompt();
+        }, /simulated inner wizard failure/);
+
+        // Flag must be reset even after failure
+        assert.strictEqual(context.suppressLoadingPrompt, false);
+    });
 });