Improve Entra account and token resolution performance

[Chris Johnstone (cjohnsto-nz)]

Description

This PR extracts the Entra authentication improvements from #22614 into a focused change.

It reduces avoidable work during Entra-backed connection setup by:

• Retrieving VS Code accounts without eagerly enumerating their tenants.
• Using an explicitly selected tenant without performing another tenant lookup.
• Resolving account IDs and labels case-insensitively, including compatible legacy account IDs.
• Avoiding eager tenant preloading when the connection dialog initializes.

These changes make Entra-backed connections faster and reduce unnecessary authentication-provider calls while preserving the existing interactive sign-in fallback.

Validation performed:

• npm run build -- --target mssql
• npm run test -- --target mssql --coverage=false
• npm run package -- --target mssql --online
• Manual VSIX installation and connection testing

Code Changes Checklist

• New or updated unit tests added
• All existing tests pass (npm run test)
• Code follows contributing guidelines
• Telemetry/logging updated if relevant — not applicable
• No regressions or UX breakage

Reviewers: Please read our reviewer guidelines

Bug: #22706

[Copilot]

Pull request overview

This PR optimizes Entra-backed authentication flows by reducing avoidable VS Code authentication-provider work during connection setup and token acquisition, primarily by deferring tenant enumeration and improving account/tenant resolution behavior.

Changes:

• Avoid eager tenant enumeration when listing VS Code accounts and when initializing the connection dialog.
• Improve account ID/label resolution (case-insensitive + legacy-compatible IDs).
• Use an explicitly selected tenant without re-enumerating tenants, and add/update unit tests for these behaviors.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
extensions/mssql/src/azure/vscodeEntraMfaUtils.ts	Updates Entra account normalization and token acquisition flow to avoid unnecessary tenant enumeration.
extensions/mssql/src/connectionconfig/azureHelpers.ts	Adjusts VS Code account lookup call sites to avoid tenant-filtering work when not needed.
extensions/mssql/src/connectionconfig/connectionDialogWebviewController.ts	Removes background preloading of tenants and avoids tenant enumeration when loading VS Code accounts.
extensions/mssql/test/unit/vscodeEntraMfaUtils.test.ts	Adds coverage for account options retrieval and token acquisition without tenant enumeration.
extensions/mssql/test/unit/connectionDialogWebviewController.test.ts	Adds a regression test ensuring tenants aren’t loaded for every account in the background.
extensions/mssql/test/unit/azureHelpers.test.ts	Updates expectations for the new getAccounts(false) usage.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Suppressed comments (4)

extensions/mssql/test/unit/vscodeEntraMfaUtils.test.ts:101

• Avoid using exact call-count assertions (like calledOnceWith) here; they’re brittle and make the test fail on unrelated refactors. Prefer asserting on the call arguments only.

        expect(getTenantsForAccount).to.have.been.calledOnceWith(mockAccounts.signedInAccount);

extensions/mssql/src/azure/vscodeEntraMfaUtils.ts:124

• tenantId is now trimmed for tenant resolution, but both error paths still pass the raw (possibly whitespace/padded) tenantId into the localized message. That can produce confusing messages that include extra whitespace. Consider trimming/normalizing the tenant ID before passing it to accountNotAvailableThroughVsCode.

    let resolvedTenantId = tenantId?.trim() || undefined;
    if (!resolvedTenantId) {
        const tenants = await VsCodeAzureHelper.getTenantsForAccount(account);
        resolvedTenantId = getDefaultTenantId(account.id, tenants);
    }

extensions/mssql/src/connectionconfig/azureHelpers.ts:163

• getAccountById is typed to always return an account, but Array.prototype.find can return undefined. This can lead to runtime errors in callers that assume a value is always returned. Consider returning Promise<… | undefined> and handling the missing-account case at call sites, or throw a clear error here.

        accountId: string,
    ): Promise<vscode.AuthenticationSessionAccountInformation> {
        const accounts = await this.getAccounts(false);
        return accounts.find((a) => a.id === accountId);
    }

extensions/mssql/src/connectionconfig/azureHelpers.ts:170

• getAccountByName has the same contract issue as getAccountById: .find(...) may return undefined, but the return type is non-optional. Align the signature/behavior to avoid surprising runtime failures.

        accountName: string,
    ): Promise<vscode.AuthenticationSessionAccountInformation> {
        const accounts = await this.getAccounts(false);
        return accounts.find((a) => a.label === accountName);
    }

[Benjin Dubishar (Benjin)]

Tested these changes against Azure SQL DBs and a Dataverse database, working well. I also did some quick perf benchmarking, and found that these changes speed up an initial connection using Entra auth by about 60%. Thanks for this PR, Chris Johnstone (@cjohnsto-nz)!

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Suppressed comments (4)

extensions/mssql/src/connectionconfig/connectionDialogWebviewController.ts:268

• Typo in comment: "intitial" should be "initial".

        // Display intitial UI since it may take a moment for the connection to load
        // due to fetching Azure account info

extensions/mssql/test/unit/connectionDialogWebviewController.test.ts:1029

• This test calls a public method via bracket notation, which bypasses type checking and is less maintainable. Since loadVscodeEntraDataAsync is public, call it directly.

            await controller["loadVscodeEntraDataAsync"]();

extensions/mssql/src/azure/vscodeEntraMfaUtils.ts:37

• The new behavior normalizes IDs to lowercase for case-insensitive compatibility, but there isn't a unit test covering mixed-case inputs. Adding a test would prevent regressions (e.g., "User@Example.com" vs "user@example.com").

    const normalizedCurrentAccountId = currentAccountId?.toLowerCase();
    const normalizedExpectedAccountId = expectedAccountId?.toLowerCase();

    return (
        !!normalizedCurrentAccountId &&
        !!normalizedExpectedAccountId &&
        (normalizedCurrentAccountId === normalizedExpectedAccountId ||
            normalizedCurrentAccountId.startsWith(normalizedExpectedAccountId) ||
            normalizedExpectedAccountId.startsWith(normalizedCurrentAccountId))

extensions/mssql/test/unit/vscodeEntraMfaUtils.test.ts:102

• Avoid brittle call-count assertions like calledOnceWith here; it tends to break on harmless refactors. It's usually enough to assert the important argument(s) were used.

        expect(result.tenantId).to.equal(expectedTenantId);
        expect(getTenantsForAccount).to.have.been.calledOnceWith(mockAccounts.signedInAccount);
    });

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Suppressed comments (4)

extensions/mssql/src/connectionconfig/azureHelpers.ts:169

• getAccountByName can return undefined when the requested account label isn’t present, but the return type claims it always returns an account. Returning undefined here can lead to downstream runtime errors; consider throwing when no match is found (or otherwise never return undefined).

        const accounts = await this.getAccounts(false);
        return accounts.find((a) => a.label === accountName);

extensions/mssql/src/connectionconfig/azureHelpers.ts:162

• getAccountById can return undefined when the requested account ID isn’t present, but the return type claims it always returns an account. Several callers immediately access .label / pass the result to tenant APIs, which can crash at runtime if the account disappears or is unavailable. Consider throwing a clear error when the account can’t be resolved (or otherwise never return undefined from this method).

This issue also appears on line 168 of the same file.

        const accounts = await this.getAccounts(false);
        return accounts.find((a) => a.id === accountId);

extensions/mssql/test/unit/vscodeEntraMfaUtils.test.ts:101

• This assertion depends on an exact call count via calledOnceWith, which is brittle (incidental extra calls will break the test). Prefer asserting on parameters without a specific count (e.g. calledWith / calledWithMatch).

        expect(getTenantsForAccount).to.have.been.calledOnceWith(mockAccounts.signedInAccount);

extensions/mssql/src/connectionconfig/connectionDialogWebviewController.ts:268

• Typo in comment: “intitial” should be “initial”.

        // Display intitial UI since it may take a moment for the connection to load
        // due to fetching Azure account info

[Chris Johnstone (cjohnsto-nz)]

Thanks Benjin Dubishar (@Benjin). I think there is still some room for improvement in the Connection Dialog too, but I'll tackle that another time 🙂

[Benjin Dubishar (Benjin)]

Chris Johnstone (@cjohnsto-nz) feel free to file some bugs/tasks with the improvements you have in mind so they don't fall through the cracks!