Improve Entra Connection Startup and Prevent Duplicate Consent Prompts

extensions/mssql/src/azure/vscodeEntraMfaUtils.ts

@@ -7,6 +7,7 @@ import * as vscode from "vscode";
 import { AzureTenant, getSessionFromVSCode } from "@microsoft/vscode-azext-azureauth";
 
 import { FormItemOptions } from "../sharedInterfaces/form";
+import { ILogger } from "../sharedInterfaces/logger";
 import { IProviderResources, IToken } from "../models/contracts/azure";
 import { getCloudProviderSettings } from "./providerSettings";
 import { VsCodeAzureHelper, getDefaultTenantId } from "../connectionconfig/azureHelpers";
@@ -19,24 +20,32 @@ export interface VscodeEntraSqlTokenInfo {
     token: IToken;
 }
 
+export interface VscodeEntraTokenTrace {
+    logger: ILogger;
+    requestId: string;
+}
+
 /**
  * Determines if the provided account IDs are compatible, meaning they are either exactly the same or one is a prefix of the other.
  */
 export function areCompatibleEntraAccountIds(
     currentAccountId?: string,
     expectedAccountId?: string,
 ): boolean {
+    const normalizedCurrentAccountId = currentAccountId?.toLowerCase();
+    const normalizedExpectedAccountId = expectedAccountId?.toLowerCase();
+
     return (
-        !!currentAccountId &&
-        !!expectedAccountId &&
-        (currentAccountId === expectedAccountId ||
-            currentAccountId.startsWith(expectedAccountId) ||
-            expectedAccountId.startsWith(currentAccountId))
+        !!normalizedCurrentAccountId &&
+        !!normalizedExpectedAccountId &&
+        (normalizedCurrentAccountId === normalizedExpectedAccountId ||
+            normalizedCurrentAccountId.startsWith(normalizedExpectedAccountId) ||
+            normalizedExpectedAccountId.startsWith(normalizedCurrentAccountId))
     );
 }
 
 export async function getVscodeEntraAccountOptions(): Promise<FormItemOptions[]> {
-    const accounts = await VsCodeAzureHelper.getAccounts();
+    const accounts = await VsCodeAzureHelper.getAccounts(false);
     return accounts.map((account) => ({
         displayName: account.label,
         value: account.id,
@@ -47,7 +56,7 @@ export async function resolveVscodeEntraAccount(
     accountId?: string,
     accountLabel?: string,
 ): Promise<vscode.AuthenticationSessionAccountInformation | undefined> {
-    const accounts = await VsCodeAzureHelper.getAccounts();
+    const accounts = await VsCodeAzureHelper.getAccounts(false);
 
     if (accountId) {
         const exactMatch = accounts.find((account) => account.id === accountId);
@@ -64,7 +73,10 @@ export async function resolveVscodeEntraAccount(
     }
 
     if (accountLabel) {
-        return accounts.find((account) => account.label === accountLabel);
+        const normalizedAccountLabel = accountLabel.trim().toLowerCase();
+        return accounts.find(
+            (account) => account.label.trim().toLowerCase() === normalizedAccountLabel,
+        );
     }
 
     return undefined;
@@ -100,8 +112,14 @@ export async function acquireTokenFromVscodeAccountForResource(
     accountId?: string,
     tenantId?: string,
     accountLabel?: string,
+    trace?: VscodeEntraTokenTrace,
 ): Promise<VscodeEntraSqlTokenInfo> {
-    const account = await resolveVscodeEntraAccount(accountId, accountLabel);
+    const account = await traceTokenPhase(
+        trace,
+        "account resolution",
+        () => resolveVscodeEntraAccount(accountId, accountLabel),
+        (resolvedAccount) => `accountResolved=${resolvedAccount !== undefined}`,
+    );
     if (!account) {
         throw new MissingEntraAuthAccountError(
             locConstants.Accounts.accountNotAvailableThroughVsCode(
@@ -111,11 +129,22 @@ export async function acquireTokenFromVscodeAccountForResource(
         );
     }
 
-    const tenants = await VsCodeAzureHelper.getTenantsForAccount(account);
-    const resolvedTenantId =
-        tenantId && tenants.some((tenant) => tenant.tenantId === tenantId)
-            ? tenantId
-            : getDefaultTenantId(account.id, tenants);
+    let resolvedTenantId = tenantId;
+    if (!resolvedTenantId) {
+        resolvedTenantId = await traceTokenPhase(
+            trace,
+            "tenant resolution",
+            async () => {
+                const tenants = await VsCodeAzureHelper.getTenantsForAccount(account);
+                return getDefaultTenantId(account.id, tenants);
+            },
+            (selectedTenantId) => `tenantResolved=${selectedTenantId !== undefined}`,
+        );
+    } else {
+        trace?.logger.debug(
+            `[ConnectionTrace] VS Code accounts tenant resolution skipped requestId=${trace.requestId} reason=explicitTenantId`,
+        );
+    }
 
     if (!resolvedTenantId) {
         throw new MissingEntraAuthAccountError(
@@ -126,16 +155,29 @@ export async function acquireTokenFromVscodeAccountForResource(
         );
     }
 
+    const silentSession = await traceTokenPhase(
+        trace,
+        "silent token session request",
+        () =>
+            getSessionFromVSCode(resourceEndpoint, resolvedTenantId, {
+                createIfNone: false,
+                silent: true,
+                account,
+            }),
+        (session) => `sessionPresent=${session !== undefined}`,
+    );
     const session =
-        (await getSessionFromVSCode(resourceEndpoint, resolvedTenantId, {
-            createIfNone: false,
-            silent: true,
-            account,
-        })) ??
-        (await getSessionFromVSCode(resourceEndpoint, resolvedTenantId, {
-            createIfNone: true,
-            account,
-        }));
+        silentSession ??
+        (await traceTokenPhase(
+            trace,
+            "interactive token session request",
+            () =>
+                getSessionFromVSCode(resourceEndpoint, resolvedTenantId, {
+                    createIfNone: true,
+                    account,
+                }),
+            (session) => `sessionPresent=${session !== undefined}`,
+        ));
 
     if (!session) {
         throw new Error(
@@ -156,6 +198,31 @@ export async function acquireTokenFromVscodeAccountForResource(
     };
 }
 
+async function traceTokenPhase<T>(
+    trace: VscodeEntraTokenTrace | undefined,
+    phase: string,
+    operation: () => Promise<T>,
+    describeResult?: (result: T) => string,
+): Promise<T> {
+    const startedAt = Date.now();
+    trace?.logger.debug(
+        `[ConnectionTrace] VS Code accounts ${phase} started requestId=${trace.requestId}`,
+    );
+
+    try {
+        const result = await operation();
+        trace?.logger.debug(
+            `[ConnectionTrace] VS Code accounts ${phase} completed requestId=${trace.requestId} durationMs=${Date.now() - startedAt}${describeResult ? ` ${describeResult(result)}` : ""}`,
+        );
+        return result;
+    } catch (error) {
+        trace?.logger.debug(
+            `[ConnectionTrace] VS Code accounts ${phase} failed requestId=${trace.requestId} durationMs=${Date.now() - startedAt}`,
+        );
+        throw error;
+    }
+}
+
 export function getCloudResourceEndpoint(endpoint: keyof IProviderResources): string {
     const cloudSettings = getCloudProviderSettings();
     const resource = cloudSettings.settings[endpoint];

extensions/mssql/src/connectionSharing/connectionSharingService.ts

@@ -49,6 +49,8 @@ export class ConnectionSharingError extends Error {
 
 export class ConnectionSharingService implements mssql.IConnectionSharingService {
     private _logger: ILogger;
+    private readonly _permissionRequests = new Map<string, Promise<boolean>>();
+
     constructor(
         private readonly _context: vscode.ExtensionContext,
         private readonly _client: SqlToolsServiceClient,
@@ -230,6 +232,26 @@ export class ConnectionSharingService implements mssql.IConnectionSharingService
     }
 
     private async requestConnectionSharingPermission(extensionId: string): Promise<boolean> {
+        const pendingRequest = this._permissionRequests.get(extensionId);
+        if (pendingRequest) {
+            return pendingRequest;
+        }
+
+        const permissionRequest = this.requestConnectionSharingPermissionCore(extensionId);
+        this._permissionRequests.set(extensionId, permissionRequest);
+
+        try {
+            return await permissionRequest;
+        } finally {
+            if (this._permissionRequests.get(extensionId) === permissionRequest) {
+                this._permissionRequests.delete(extensionId);
+            }
+        }
+    }
+
+    private async requestConnectionSharingPermissionCore(
+        extensionId: string,
+    ): Promise<boolean> {
         this._logger.info(`Requesting connection sharing permission for extension: ${extensionId}`);
 
         const currentPermission = await this.getExtensionPermission(extensionId);

extensions/mssql/src/connectionconfig/azureHelpers.ts

@@ -158,14 +158,14 @@ export class VsCodeAzureHelper {
     public static async getAccountById(
         accountId: string,
     ): Promise<vscode.AuthenticationSessionAccountInformation> {
-        const accounts = await this.getAccounts();
+        const accounts = await this.getAccounts(false);
         return accounts.find((a) => a.id === accountId);
     }
 
     public static async getAccountByName(
         accountName: string,
     ): Promise<vscode.AuthenticationSessionAccountInformation> {
-        const accounts = await this.getAccounts();
+        const accounts = await this.getAccounts(false);
         return accounts.find((a) => a.label === accountName);
     }
 
@@ -190,7 +190,7 @@ export class VsCodeAzureHelper {
 
         if (forceSignInPrompt || !(await auth.isSignedIn())) {
             const accountsBefore = new Set(
-                (await VsCodeAzureHelper.getAccounts()).map((a) => a.id),
+                (await VsCodeAzureHelper.getAccounts(false)).map((a) => a.id),
             );
 
             const result = await auth.signIn();
@@ -199,14 +199,14 @@ export class VsCodeAzureHelper {
                 throw new Error("Azure sign-in was canceled or failed.");
             }
 
-            const accountsAfter = await VsCodeAzureHelper.getAccounts();
+            const accountsAfter = await VsCodeAzureHelper.getAccounts(false);
             const newAccount = accountsAfter.find((a) => !accountsBefore.has(a.id));
 
             return { auth, newAccountId: newAccount?.id ?? accountsAfter[0]?.id };
         }
 
         // Already signed in — return the first available account
-        const accounts = await VsCodeAzureHelper.getAccounts();
+        const accounts = await VsCodeAzureHelper.getAccounts(false);
         return { auth, newAccountId: accounts[0]?.id };
     }