Upgrade to auth v6 and improve extension performance

package.json

@@ -7,7 +7,7 @@
     "icon": "resources/resourceGroup.png",
     "aiKey": "0c6ae279ed8443289764825290e4f9e2-1a736e7c-1324-4338-be46-fc2a58ae4d14-7255",
     "engines": {
-        "vscode": "^1.105.0"
+        "vscode": "^1.106.0"
     },
     "repository": {
         "type": "git",
@@ -926,7 +926,7 @@
     "dependencies": {
         "@azure/arm-resources": "^5.2.0",
         "@azure/arm-resources-profile-2020-09-01-hybrid": "^2.1.0",
-        "@microsoft/vscode-azext-azureauth": "^5.1.1",
+        "@microsoft/vscode-azext-azureauth": "^6.0.0-alpha.1",
         "@microsoft/vscode-azext-azureutils": "^4.0.0",
         "@microsoft/vscode-azext-utils": "^4.0.4",
         "form-data": "^4.0.4",

src/cloudConsole/cloudConsole.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { TenantIdDescription } from '@azure/arm-resources-subscriptions';
-import { AzureSubscriptionProvider, getConfiguredAzureEnv } from '@microsoft/vscode-azext-azureauth';
+import { AzureSubscriptionProvider, AzureTenant, getConfiguredAzureEnv } from '@microsoft/vscode-azext-azureauth';
 import { IActionContext, IAzureQuickPickItem, IParsedError, callWithTelemetryAndErrorHandlingSync, nonNullProp, parseError } from '@microsoft/vscode-azext-utils';
 import * as cp from 'child_process';
 import { default as FormData } from 'form-data';
@@ -262,20 +262,18 @@ export function createCloudConsole(subscriptionProvider: AzureSubscriptionProvid
                 res.end();
             });
 
-            if (!await subscriptionProvider.isSignedIn()) {
-                serverQueue.push({ type: 'log', args: [localize('loggingIn', "Signing in...")] });
-                try {
-                    if (await subscriptionProvider.signIn()) {
-                        serverQueue.push({ type: 'log', args: [localize('loggingIn', "Signed in successful.")] });
-                    }
-                } catch (e) {
-                    serverQueue.push({ type: 'log', args: [localize('loggingIn', parseError(e).message)] });
-                    // We used to delay for a second then exit here, but then the user can't read or copy the error message
-                    // await delay(1000);
-                    // serverQueue.push({ type: 'exit' });
-                    updateStatus('Disconnected');
-                    return;
+            serverQueue.push({ type: 'log', args: [localize('loggingIn', "Signing in...")] });
+            try {
+                if (await subscriptionProvider.signIn()) {
+                    serverQueue.push({ type: 'log', args: [localize('loggingIn', "Signed in successful.")] });
                 }
+            } catch (e) {
+                serverQueue.push({ type: 'log', args: [localize('loggingIn', parseError(e).message)] });
+                // We used to delay for a second then exit here, but then the user can't read or copy the error message
+                // await delay(1000);
+                // serverQueue.push({ type: 'exit' });
+                updateStatus('Disconnected');
+                return;
             }
 
             const env: TerminalOptions['env'] = {
@@ -336,10 +334,13 @@ export function createCloudConsole(subscriptionProvider: AzureSubscriptionProvid
 
             liveServerQueue = serverQueue;
 
-            const tenants = await subscriptionProvider.getTenants();
+            const tenants: AzureTenant[] = [];
+            for (const account of await subscriptionProvider.getAccounts({ filter: false })) {
+                tenants.push(...await subscriptionProvider.getTenantsForAccount(account, { filter: false }));
+            }
             let selectedTenant: TenantIdDescription | undefined = undefined;
 
-            const subscriptions = await subscriptionProvider.getSubscriptions(false);
+            const subscriptions = await subscriptionProvider.getAvailableSubscriptions({ filter: false });
             if (tenants.length <= 1) {
                 serverQueue.push({ type: 'log', args: [localize('foundOneTenant', `Found 1 tenant.`)] });
                 // if they have only one tenant, use it

src/commands/accounts/logIn.ts

@@ -18,6 +18,8 @@ export async function logIn(_context: IActionContext): Promise<void> {
         await provider.signIn();
     } finally {
         _isLoggingIn = false;
+        // Clear cache to ensure fresh data is fetched after sign-in
+        ext.setClearCacheOnNextLoad();
         ext.actions.refreshAzureTree(); // Refresh now that sign in is complete
         ext.actions.refreshTenantTree(); // Refresh now that sign in is complete
     }

src/commands/accounts/selectSubscriptions.ts

@@ -3,7 +3,7 @@
 *  Licensed under the MIT License. See License.md in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
 
-import { AzureSubscription } from "@microsoft/vscode-azext-azureauth";
+import { AzureSubscription, isNotSignedInError } from "@microsoft/vscode-azext-azureauth";
 import { IActionContext, IAzureQuickPickItem, nonNullValue } from "@microsoft/vscode-azext-utils";
 import * as vscode from "vscode";
 import { ext } from "../../extensionVariables";
@@ -28,16 +28,16 @@ export interface SelectSubscriptionOptions {
 
 export async function selectSubscriptions(context: IActionContext, options?: SelectSubscriptionOptions): Promise<void> {
     const provider = await ext.subscriptionProviderFactory();
-    if (await provider.isSignedIn()) {
-
+    try {
         const selectedSubscriptionsWithFullId = await getSelectedTenantAndSubscriptionIds();
         const selectedSubscriptionIds = selectedSubscriptionsWithFullId.map(id => id.split('/')[1]);
 
         let subscriptionsShownInPicker: string[] = [];
 
         const subscriptionQuickPickItems: () => Promise<IAzureQuickPickItem<AzureSubscription>[]> = async () => {
-            // If there are no tenants selected by default all subscriptions will be shown.
-            let subscriptions = await provider.getSubscriptions(false);
+            // Use cached subscriptions if available - the cache is cleared on sign-in/sign-out
+            // This avoids unnecessary network requests when the user opens the subscription picker
+            let subscriptions = await provider.getAvailableSubscriptions({ filter: false });
 
             if (options?.account || options?.tenantId) {
 
@@ -89,13 +89,17 @@ export async function selectSubscriptions(context: IActionContext, options?: Sel
         }
 
         ext.actions.refreshAzureTree();
-    } else {
-        const signIn: vscode.MessageItem = { title: localize('signIn', 'Sign In') };
-        void vscode.window.showInformationMessage(localize('notSignedIn', 'You are not signed in. Sign in to continue.'), signIn).then((input) => {
-            if (input === signIn) {
-                void provider.signIn();
-            }
-        });
+    } catch (error) {
+        if (isNotSignedInError(error)) {
+            const signIn: vscode.MessageItem = { title: localize('signIn', 'Sign In') };
+            void vscode.window.showInformationMessage(localize('notSignedIn', 'You are not signed in. Sign in to continue.'), signIn).then((input) => {
+                if (input === signIn) {
+                    void provider.signIn();
+                }
+            });
+        } else {
+            throw error;
+        }
     }
 }
 

src/commands/registerCommands.ts

@@ -41,10 +41,19 @@ export function registerCommands(): void {
     registerCommand('azureResourceGroups.uploadFileCloudConsole', uploadFileToCloudShell);
 
     // Special-case refresh that ignores the selected/focused node and always refreshes the entire tree. Used by the refresh button in the tree title.
-    registerCommand('azureResourceGroups.refreshTree', () => ext.actions.refreshAzureTree());
+    registerCommand('azureResourceGroups.refreshTree', () => {
+        ext.setClearCacheOnNextLoad();
+        ext.actions.refreshAzureTree();
+    });
     registerCommand('azureWorkspace.refreshTree', () => ext.actions.refreshWorkspaceTree());
-    registerCommand('azureFocusView.refreshTree', () => ext.actions.refreshFocusTree());
-    registerCommand('azureTenantsView.refreshTree', () => ext.actions.refreshTenantTree());
+    registerCommand('azureFocusView.refreshTree', () => {
+        ext.setClearCacheOnNextLoad();
+        ext.actions.refreshFocusTree();
+    });
+    registerCommand('azureTenantsView.refreshTree', () => {
+        ext.setClearCacheOnNextLoad();
+        ext.actions.refreshTenantTree();
+    });
 
     // v1.5 client extensions attach these commands to tree item context menus for refreshing their tree items
     registerCommand('azureResourceGroups.refresh', async (context, node?: ResourceGroupsItem) => {
@@ -78,8 +87,10 @@ export function registerCommands(): void {
     });
 
     registerCommand('azureTenantsView.signInToTenant', async (_context, node: TenantTreeItem) => {
-        await (await ext.subscriptionProviderFactory()).signIn(node.tenantId, node.account);
-        ext.actions.refreshTenantTree(node);
+        await (await ext.subscriptionProviderFactory()).signIn(node);
+        ext.setClearCacheOnNextLoad();
+        ext.actions.refreshTenantTree();
+        ext.actions.refreshAzureTree();
     });
 
     registerCommand('azureResourceGroups.focusGroup', focusGroup);
@@ -88,7 +99,12 @@ export function registerCommands(): void {
     registerCommand('azureResourceGroups.logIn', (context: IActionContext) => logIn(context));
     registerCommand('azureTenantsView.addAccount', (context: IActionContext) => logIn(context));
     registerCommand('azureResourceGroups.selectSubscriptions', (context: IActionContext, options: SelectSubscriptionOptions) => selectSubscriptions(context, options));
-    registerCommand('azureResourceGroups.signInToTenant', async () => signInToTenant(await ext.subscriptionProviderFactory()));
+    registerCommand('azureResourceGroups.signInToTenant', async () => {
+        await signInToTenant(await ext.subscriptionProviderFactory());
+        ext.setClearCacheOnNextLoad();
+        ext.actions.refreshTenantTree();
+        ext.actions.refreshAzureTree();
+    });
 
     registerCommand('azureResourceGroups.createResourceGroup', createResourceGroup);
     registerCommand('azureResourceGroups.deleteResourceGroupV2', deleteResourceGroupV2);

src/commands/sovereignCloud/configureSovereignCloud.ts

@@ -26,7 +26,9 @@ export async function configureSovereignCloud(context: ConfigureSovereignCloudCo
     await wizard.prompt();
     await wizard.execute();
 
-    // refresh resources and tenant view to accurrately reflect information for the selected sovereign cloud
+    // Clear cache and refresh views to reflect the selected sovereign cloud
+    // This ensures accounts from the previous environment are not shown
+    ext.setClearCacheOnNextLoad();
     ext.actions.refreshAzureTree();
     ext.actions.refreshTenantTree();
 }

src/extension.ts

@@ -91,7 +91,7 @@ export async function activate(context: vscode.ExtensionContext, perfStats: { lo
         activateContext.telemetry.measurements.mainFileLoad = (perfStats.loadEndTime - perfStats.loadStartTime) / 1000;
 
 
-        ext.subscriptionProviderFactory = getSubscriptionProviderFactory(activateContext);
+        ext.subscriptionProviderFactory = getSubscriptionProviderFactory();
 
         ext.tagFS = new TagFileSystem(ext.appResourceTree);
         context.subscriptions.push(vscode.workspace.registerFileSystemProvider(TagFileSystem.scheme, ext.tagFS));
@@ -213,7 +213,7 @@ export async function activate(context: vscode.ExtensionContext, perfStats: { lo
     ext.workspaceTree = new CompatibleAzExtTreeDataProvider(workspaceResourceTreeDataProvider);
 
     const getSubscriptions: (filter: boolean) => Promise<AzureSubscription[]> =
-        async (filter: boolean) => { return await (await ext.subscriptionProviderFactory()).getSubscriptions(filter); };
+        async (filter: boolean) => { return await (await ext.subscriptionProviderFactory()).getAvailableSubscriptions({ filter }); };
 
     const coreApiFactories: AzureExtensionApiFactory[] = [
         {

src/extensionVariables.ts

@@ -53,6 +53,30 @@ export namespace ext {
     export let subscriptionProviderFactory: () => Promise<AzureSubscriptionProvider>;
     export let managedIdentityBranchDataProvider: ManagedIdentityBranchDataProvider;
 
+    /**
+     * Cache invalidation flag. When set to true, the next call to `consumeClearCacheFlag()`
+     * will return true and atomically reset the flag to false. This prevents race conditions
+     * where multiple trees might read and reset the flag independently.
+     */
+    let clearCacheOnNextLoadFlag: boolean = false;
+
+    /**
+     * Sets the flag to clear auth caches on the next load.
+     */
+    export function setClearCacheOnNextLoad(): void {
+        clearCacheOnNextLoadFlag = true;
+    }
+
+    /**
+     * Atomically consumes the clear cache flag. Returns true if caches should be cleared,
+     * and resets the flag to false. This ensures only the first consumer gets `true`.
+     */
+    export function consumeClearCacheFlag(): boolean {
+        const shouldClear = clearCacheOnNextLoadFlag;
+        clearCacheOnNextLoadFlag = false;
+        return shouldClear;
+    }
+
     // eslint-disable-next-line @typescript-eslint/no-namespace
     export namespace v2 {
         export let api: AzureResourcesApiInternal;

src/managedIdentity/TargetServiceRoleAssignmentItem.ts

@@ -29,7 +29,7 @@ export class TargetServiceRoleAssignmentItem implements TreeElementBase {
 
             if (this._loadedAllSubscriptions) {
                 // filter out this sub since it's already loaded
-                const subscriptions = (await (await ext.subscriptionProviderFactory()).getSubscriptions(false)).filter(s => s.subscriptionId !== this.subscription.subscriptionId);
+                const subscriptions = (await (await ext.subscriptionProviderFactory()).getAvailableSubscriptions({ filter: false })).filter(s => s.subscriptionId !== this.subscription.subscriptionId);
                 await Promise.allSettled(subscriptions.map(async (subscription) => {
                     children.push(...await createRoleDefinitionsItems(context, subscription, this.msi, this.subscription.subscriptionId));
                 }));

src/services/getSubscriptionProviderFactory.ts

@@ -1,47 +1,17 @@
 /*---------------------------------------------------------------------------------------------
  *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *  Licensed under the MIT License. See License.md in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { AzureDevOpsSubscriptionProviderInitializer, AzureSubscriptionProvider, createAzureDevOpsSubscriptionProviderFactory } from "@microsoft/vscode-azext-azureauth";
-import { IActionContext } from "@microsoft/vscode-azext-utils";
+import type { AzureSubscriptionProvider } from "@microsoft/vscode-azext-azureauth";
 import { createVSCodeAzureSubscriptionProviderFactory } from "./VSCodeAzureSubscriptionProvider";
 
 /**
  * Returns a factory function that creates a subscription provider, satisfying the `AzureSubscriptionProvider` interface.
  *
- * If the `useAzureSubscriptionProvider` is set to `true`, an `AzureDevOpsSubscriptionProviderFactory` is returned.
- * Otherwise, a `VSCodeSubscriptionProviderFactory` is returned.
- *
+ * For nightly tests that require Azure DevOps federated credentials, use the test API to set
+ * `ext.testing.overrideAzureSubscriptionProvider` with an AzDO provider factory instead.
  */
-export function getSubscriptionProviderFactory(activateContext?: IActionContext): () => Promise<AzureSubscriptionProvider> {
-    // if this for a nightly test, we want to use the test subscription provider
-    const useAzureFederatedCredentials: boolean = !/^(false|0)?$/i.test(process.env['AzCode_UseAzureFederatedCredentials'] || '');
-    if (useAzureFederatedCredentials) {
-        // when running tests, ensure we throw the errors and they aren't silently swallowed
-        if (activateContext) {
-            activateContext.errorHandling.rethrow = useAzureFederatedCredentials;
-        }
-
-        const serviceConnectionId: string | undefined = process.env['AzCode_ServiceConnectionID'];
-        const domain: string | undefined = process.env['AzCode_ServiceConnectionDomain'];
-        const clientId: string | undefined = process.env['AzCode_ServiceConnectionClientID'];
-
-        if (!serviceConnectionId || !domain || !clientId) {
-            throw new Error(`Using Azure DevOps federated credentials, but federated service connection is not configured\n
-                                process.env.AzCode_ServiceConnectionID: ${serviceConnectionId ? "✅" : "❌"}\n
-                                process.env.AzCode_ServiceConnectionDomain: ${domain ? "✅" : "❌"}\n
-                                process.env.AzCode_ServiceConnectionClientID: ${clientId ? "✅" : "❌"}\n
-                            `);
-        }
-
-        const initializer: AzureDevOpsSubscriptionProviderInitializer = {
-            serviceConnectionId,
-            domain,
-            clientId,
-        };
-        return createAzureDevOpsSubscriptionProviderFactory(initializer);
-    } else {
-        return createVSCodeAzureSubscriptionProviderFactory();
-    }
+export function getSubscriptionProviderFactory(): () => Promise<AzureSubscriptionProvider> {
+    return createVSCodeAzureSubscriptionProviderFactory();
 }