Fix the 'Add account' for AlwaysEncrypted sign-in (#21432)

* Fix the 'Add account' for AlwaysEncrypted sign-in

* Add a function comment for the null/undefined behavior

Author: kburtram

extensions/mssql/src/controllers/connectionManager.ts

@@ -2036,33 +2036,43 @@ export default class ConnectionManager {
     private async handleSecurityTokenRequest(
         params: RequestSecurityTokenParams,
     ): Promise<RequestSecurityTokenResponse> {
-        if (this._keyVaultTokenCache.has(JSON.stringify(params))) {
-            const token = this._keyVaultTokenCache.get(JSON.stringify(params));
-            const isExpired = AzureController.isTokenExpired(token.expiresOn);
-            if (!isExpired) {
-                return {
-                    accountKey: token.key,
-                    token: token.token,
-                };
-            } else {
-                this._keyVaultTokenCache.delete(JSON.stringify(params));
+        try {
+            if (this._keyVaultTokenCache.has(JSON.stringify(params))) {
+                const token = this._keyVaultTokenCache.get(JSON.stringify(params));
+                const isExpired = AzureController.isTokenExpired(token.expiresOn);
+                if (!isExpired) {
+                    return {
+                        accountKey: token.key,
+                        token: token.token,
+                    };
+                } else {
+                    this._keyVaultTokenCache.delete(JSON.stringify(params));
+                }
             }
-        }
-        const account = await this.selectAccount();
-        const tenant = await this.selectTenantId(account);
+            const account = await this.selectAccount();
+            const tenant = await this.selectTenantId(account);
 
-        const token = await this.azureController.getAccountSecurityToken(
-            account,
-            tenant,
-            getCloudProviderSettings(account.key.providerId).settings.azureKeyVaultResource,
-        );
+            const token = await this.azureController.getAccountSecurityToken(
+                account,
+                tenant,
+                getCloudProviderSettings(account.key.providerId).settings.azureKeyVaultResource,
+            );
 
-        this._keyVaultTokenCache.set(JSON.stringify(params), token);
+            this._keyVaultTokenCache.set(JSON.stringify(params), token);
 
-        return {
-            accountKey: token.key,
-            token: token.token,
-        };
+            return {
+                accountKey: token.key,
+                token: token.token,
+            };
+        } catch (error) {
+            this._logger.error(`Security token request failed: ${getErrorMessage(error)}`);
+            // Return empty response rather than letting the error propagate
+            // to STS as a null reference
+            return {
+                accountKey: "",
+                token: "",
+            };
+        }
     }
 
     private async selectAccount(): Promise<IAccount> {
@@ -2074,6 +2084,16 @@ export default class ConnectionManager {
         const quickPickItems = this.createAccountQuickPickItems(accounts, currentAccountId);
         const selectedAccount = await this.showAccountQuickPick(quickPickItems);
 
+        // eslint-disable-next-line no-restricted-syntax
+        if (selectedAccount === null) {
+            // User selected "Sign in to Azure" — trigger sign-in and use the new account
+            const newAccount = await this.addAccount();
+            if (!newAccount) {
+                throw new Error(LocalizedConstants.Connection.noAccountSelected);
+            }
+            return newAccount;
+        }
+
         if (!selectedAccount) {
             throw new Error(LocalizedConstants.Connection.noAccountSelected);
         }
@@ -2103,30 +2123,47 @@ export default class ConnectionManager {
         return accountItems;
     }
 
+    /*
+     * Shows a quick pick to select an account. Returns the selected account, null if "Sign in to Azure" was selected,
+     * or undefined if the quick pick was dismissed.
+     * @params items The quick pick items to show
+     * @returns The selected account, null if "Sign in to Azure" was selected, or undefined if the quick pick was dismissed
+     */
     private async showAccountQuickPick(
         items: AccountQuickPickItem[],
-    ): Promise<IAccount | undefined> {
-        const account = await new Promise<IAccount | undefined>((resolve, reject) => {
+    ): Promise<IAccount | null | undefined> {
+        const account = await new Promise<IAccount | null | undefined>((resolve, reject) => {
             const quickPick = vscode.window.createQuickPick<AccountQuickPickItem>();
             quickPick.items = items;
             quickPick.placeholder = LocalizedConstants.Connection.SelectAccountForKeyVault;
+            let accepted = false;
 
             quickPick.onDidAccept(async () => {
                 try {
+                    accepted = true;
                     const selectedItem = quickPick.selectedItems[0];
+                    quickPick.dispose();
                     if (!selectedItem) {
                         resolve(undefined);
                         return;
                     }
 
                     const account = selectedItem.account;
-                    quickPick.dispose();
-                    resolve(account);
+                    // Return null to signal "sign in" was selected (account is undefined on that item)
+                    resolve(account ?? null); // eslint-disable-line no-restricted-syntax
                 } catch (error) {
                     quickPick.dispose();
                     reject(error);
                 }
             });
+
+            quickPick.onDidHide(() => {
+                quickPick.dispose();
+                if (!accepted) {
+                    resolve(undefined);
+                }
+            });
+
             quickPick.show();
         });
         return account;