Better lifecycle handling (#242758)

I moved to a factory model because there was just so much that needed to be async.

I think the amount of async code will be reduced in the future as we remove some migration logic, but this makes sure we don't accidentally create instances without awaiting their initialization.

extensions/microsoft-authentication/src/common/accountAccess.ts

@@ -7,32 +7,47 @@ import { Disposable, Event, EventEmitter, LogOutputChannel, SecretStorage } from
 import { AccountInfo } from '@azure/msal-node';
 
 export interface IAccountAccess {
-	initialize(): Promise<void>;
 	onDidAccountAccessChange: Event<void>;
 	isAllowedAccess(account: AccountInfo): boolean;
 	setAllowedAccess(account: AccountInfo, allowed: boolean): Promise<void>;
 }
 
-export class ScopedAccountAccess implements IAccountAccess {
+export class ScopedAccountAccess implements IAccountAccess, Disposable {
 	private readonly _onDidAccountAccessChangeEmitter = new EventEmitter<void>();
 	readonly onDidAccountAccessChange = this._onDidAccountAccessChangeEmitter.event;
 
-	private readonly _accountAccessSecretStorage: AccountAccessSecretStorage;
-
 	private value = new Array<string>();
 
-	constructor(
+	private readonly _disposable: Disposable;
+
+	private constructor(
+		private readonly _accountAccessSecretStorage: IAccountAccessSecretStorage,
+		disposables: Disposable[] = []
+	) {
+		this._disposable = Disposable.from(
+			...disposables,
+			this._onDidAccountAccessChangeEmitter,
+			this._accountAccessSecretStorage.onDidChange(() => this.update())
+		);
+	}
+
+	static async create(
 		secretStorage: SecretStorage,
 		cloudName: string,
 		logger: LogOutputChannel,
-		migrations?: { clientId: string; authority: string }[],
-	) {
-		this._accountAccessSecretStorage = new AccountAccessSecretStorage(secretStorage, cloudName, logger, migrations);
-		this._accountAccessSecretStorage.onDidChange(() => this.update());
+		migrations: { clientId: string; authority: string }[] | undefined,
+	): Promise<ScopedAccountAccess> {
+		const storage = await AccountAccessSecretStorage.create(secretStorage, cloudName, logger, migrations);
+		const access = new ScopedAccountAccess(storage, [storage]);
+		await access.initialize();
+		return access;
 	}
 
-	async initialize() {
-		await this._accountAccessSecretStorage.initialize();
+	dispose() {
+		this._disposable.dispose();
+	}
+
+	private async initialize(): Promise<void> {
 		await this.update();
 	}
 
@@ -62,15 +77,22 @@ export class ScopedAccountAccess implements IAccountAccess {
 	}
 }
 
-export class AccountAccessSecretStorage {
+interface IAccountAccessSecretStorage {
+	get(): Promise<string[] | undefined>;
+	store(value: string[]): Thenable<void>;
+	delete(): Thenable<void>;
+	onDidChange: Event<void>;
+}
+
+class AccountAccessSecretStorage implements IAccountAccessSecretStorage, Disposable {
 	private _disposable: Disposable;
 
-	private readonly _onDidChangeEmitter = new EventEmitter<void>;
+	private readonly _onDidChangeEmitter = new EventEmitter<void>();
 	readonly onDidChange: Event<void> = this._onDidChangeEmitter.event;
 
 	private readonly _key = `accounts-${this._cloudName}`;
 
-	constructor(
+	private constructor(
 		private readonly _secretStorage: SecretStorage,
 		private readonly _cloudName: string,
 		private readonly _logger: LogOutputChannel,
@@ -86,10 +108,21 @@ export class AccountAccessSecretStorage {
 		);
 	}
 
+	static async create(
+		secretStorage: SecretStorage,
+		cloudName: string,
+		logger: LogOutputChannel,
+		migrations?: { clientId: string; authority: string }[],
+	): Promise<AccountAccessSecretStorage> {
+		const storage = new AccountAccessSecretStorage(secretStorage, cloudName, logger, migrations);
+		await storage.initialize();
+		return storage;
+	}
+
 	/**
 	 * TODO: Remove this method after a release with the migration
 	 */
-	async initialize(): Promise<void> {
+	private async initialize(): Promise<void> {
 		if (!this._migrations) {
 			return;
 		}

extensions/microsoft-authentication/src/common/cachePlugin.ts

@@ -6,7 +6,7 @@
 import { ICachePlugin, TokenCacheContext } from '@azure/msal-node';
 import { Disposable, EventEmitter, SecretStorage } from 'vscode';
 
-export class SecretStorageCachePlugin implements ICachePlugin {
+export class SecretStorageCachePlugin implements ICachePlugin, Disposable {
 	private readonly _onDidChange: EventEmitter<void> = new EventEmitter<void>();
 	readonly onDidChange = this._onDidChange.event;
 

extensions/microsoft-authentication/src/common/publicClientCache.ts

@@ -5,8 +5,7 @@
 import type { AccountInfo, AuthenticationResult, InteractiveRequest, RefreshTokenRequest, SilentFlowRequest } from '@azure/msal-node';
 import type { Disposable, Event } from 'vscode';
 
-export interface ICachedPublicClientApplication extends Disposable {
-	initialize(): Promise<void>;
+export interface ICachedPublicClientApplication {
 	onDidAccountsChange: Event<{ added: AccountInfo[]; changed: AccountInfo[]; deleted: AccountInfo[] }>;
 	onDidRemoveLastAccount: Event<void>;
 	acquireTokenSilent(request: SilentFlowRequest): Promise<AuthenticationResult>;

extensions/microsoft-authentication/src/extensionV2.ts

@@ -49,14 +49,13 @@ async function initMicrosoftSovereignCloudAuthProvider(
 		return undefined;
 	}
 
-	const authProvider = new MsalAuthProvider(
+	const authProvider = await MsalAuthProvider.create(
 		context,
 		new MicrosoftSovereignCloudAuthenticationTelemetryReporter(context.extension.packageJSON.aiKey),
 		window.createOutputChannel(l10n.t('Microsoft Sovereign Cloud Authentication'), { log: true }),
 		uriHandler,
 		env
 	);
-	await authProvider.initialize();
 	const disposable = authentication.registerAuthenticationProvider(
 		'microsoft-sovereign-cloud',
 		authProviderName,
@@ -70,13 +69,12 @@ async function initMicrosoftSovereignCloudAuthProvider(
 export async function activate(context: ExtensionContext, mainTelemetryReporter: MicrosoftAuthenticationTelemetryReporter) {
 	const uriHandler = new UriEventHandler();
 	context.subscriptions.push(uriHandler);
-	const authProvider = new MsalAuthProvider(
+	const authProvider = await MsalAuthProvider.create(
 		context,
 		mainTelemetryReporter,
 		Logger,
 		uriHandler
 	);
-	await authProvider.initialize();
 	context.subscriptions.push(authentication.registerAuthenticationProvider(
 		'microsoft',
 		'Microsoft',

extensions/microsoft-authentication/src/node/authProvider.ts

@@ -7,7 +7,7 @@ import { AuthenticationGetSessionOptions, AuthenticationProvider, Authentication
 import { Environment } from '@azure/ms-rest-azure-env';
 import { CachedPublicClientApplicationManager } from './publicClientCache';
 import { UriEventHandler } from '../UriEventHandler';
-import { ICachedPublicClientApplication } from '../common/publicClientCache';
+import { ICachedPublicClientApplication, ICachedPublicClientApplicationManager } from '../common/publicClientCache';
 import { MicrosoftAccountType, MicrosoftAuthenticationTelemetryReporter } from '../common/telemetryReporter';
 import { ScopeData } from '../common/scopeData';
 import { EventBufferer } from '../common/event';
@@ -22,7 +22,6 @@ const MSA_PASSTHRU_TID = 'f8cdef31-a31e-4b4a-93e4-5f571e91255a';
 export class MsalAuthProvider implements AuthenticationProvider {
 
 	private readonly _disposables: { dispose(): void }[];
-	private readonly _publicClientManager: CachedPublicClientApplicationManager;
 	private readonly _eventBufferer = new EventBufferer();
 
 	/**
@@ -43,15 +42,15 @@ export class MsalAuthProvider implements AuthenticationProvider {
 	 */
 	onDidChangeSessions = this._onDidChangeSessionsEmitter.event;
 
-	constructor(
+	private constructor(
 		private readonly _context: ExtensionContext,
 		private readonly _telemetryReporter: MicrosoftAuthenticationTelemetryReporter,
 		private readonly _logger: LogOutputChannel,
 		private readonly _uriHandler: UriEventHandler,
+		private readonly _publicClientManager: ICachedPublicClientApplicationManager,
 		private readonly _env: Environment = Environment.AzureCloud
 	) {
 		this._disposables = _context.subscriptions;
-		this._publicClientManager = new CachedPublicClientApplicationManager(_context.secrets, this._logger, this._env.name);
 		const accountChangeEvent = this._eventBufferer.wrapEvent(
 			this._publicClientManager.onDidAccountsChange,
 			(last, newEvent) => {
@@ -76,11 +75,24 @@ export class MsalAuthProvider implements AuthenticationProvider {
 		)(e => this._handleAccountChange(e));
 		this._disposables.push(
 			this._onDidChangeSessionsEmitter,
-			this._publicClientManager,
 			accountChangeEvent
 		);
 	}
 
+	static async create(
+		context: ExtensionContext,
+		telemetryReporter: MicrosoftAuthenticationTelemetryReporter,
+		logger: LogOutputChannel,
+		uriHandler: UriEventHandler,
+		env: Environment = Environment.AzureCloud
+	): Promise<MsalAuthProvider> {
+		const publicClientManager = await CachedPublicClientApplicationManager.create(context.secrets, logger, env.name);
+		context.subscriptions.push(publicClientManager);
+		const authProvider = new MsalAuthProvider(context, telemetryReporter, logger, uriHandler, publicClientManager, env);
+		await authProvider.initialize();
+		return authProvider;
+	}
+
 	/**
 	 * Migrate sessions from the old secret storage to MSAL.
 	 * TODO: MSAL Migration. Remove this when we remove the old flow.
@@ -109,9 +121,7 @@ export class MsalAuthProvider implements AuthenticationProvider {
 		}
 	}
 
-	async initialize(): Promise<void> {
-		await this._eventBufferer.bufferEventsAsync(() => this._publicClientManager.initialize());
-
+	private async initialize(): Promise<void> {
 		if (!this._context.globalState.get('msalMigration', false)) {
 			await this._migrateSessions();
 		}

extensions/microsoft-authentication/src/node/cachedPublicClientApplication.ts

@@ -39,7 +39,7 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 
 	//#endregion
 
-	constructor(
+	private constructor(
 		private readonly _clientId: string,
 		private readonly _secretStorage: SecretStorage,
 		private readonly _accountAccess: IAccountAccess,
@@ -47,7 +47,7 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 	) {
 		const loggerOptions = new MsalLoggerOptions(_logger);
 		const nativeBrokerPlugin = new NativeBrokerPlugin();
-		this._isBrokerAvailable = nativeBrokerPlugin.isBrokerAvailable ?? false;
+		this._isBrokerAvailable = nativeBrokerPlugin.isBrokerAvailable;
 		this._pca = new PublicClientApplication({
 			auth: { clientId: _clientId },
 			system: {
@@ -63,17 +63,26 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 		this._disposable = Disposable.from(
 			this._registerOnSecretStorageChanged(),
 			this._onDidAccountsChangeEmitter,
-			this._onDidRemoveLastAccountEmitter
+			this._onDidRemoveLastAccountEmitter,
+			this._secretStorageCachePlugin
 		);
 	}
 
 	get accounts(): AccountInfo[] { return this._accounts; }
 	get clientId(): string { return this._clientId; }
 
-	async initialize(): Promise<void> {
-		if (this._isBrokerAvailable) {
-			await this._accountAccess.initialize();
-		}
+	static async create(
+		clientId: string,
+		secretStorage: SecretStorage,
+		accountAccess: IAccountAccess,
+		logger: LogOutputChannel
+	): Promise<CachedPublicClientApplication> {
+		const app = new CachedPublicClientApplication(clientId, secretStorage, accountAccess, logger);
+		await app.initialize();
+		return app;
+	}
+
+	private async initialize(): Promise<void> {
 		await this._sequencer.queue(() => this._update());
 	}
 
@@ -178,7 +187,15 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 	 */
 	async acquireTokenByRefreshToken(request: RefreshTokenRequest): Promise<AuthenticationResult | null> {
 		this._logger.debug(`[acquireTokenByRefreshToken] [${this._clientId}] [${request.authority}] [${request.scopes.join(' ')}]`);
-		const result = await this._sequencer.queue(() => this._pca.acquireTokenByRefreshToken(request));
+		const result = await this._sequencer.queue(async () => {
+			const result = await this._pca.acquireTokenByRefreshToken(request);
+			// Force an update so that the account cache is updated.
+			// TODO:@TylerLeonhardt The problem is, we use the sequencer for
+			// change events but we _don't_ use it for the accounts cache.
+			// We should probably use it for the accounts cache as well.
+			await this._update();
+			return result;
+		});
 		if (result) {
 			// this._setupRefresh(result);
 			if (this._isBrokerAvailable && result.account) {

extensions/microsoft-authentication/src/node/publicClientCache.ts

@@ -20,37 +20,46 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient
 	private readonly _pcaDisposables = new Map<string, Disposable>();
 
 	private _disposable: Disposable;
-	private _pcasSecretStorage: PublicClientApplicationsSecretStorage;
-	private _accountAccess: IAccountAccess | undefined;
 
 	private readonly _onDidAccountsChangeEmitter = new EventEmitter<{ added: AccountInfo[]; changed: AccountInfo[]; deleted: AccountInfo[] }>();
 	readonly onDidAccountsChange = this._onDidAccountsChangeEmitter.event;
 
-	constructor(
+	private constructor(
+		private readonly _pcasSecretStorage: IPublicClientApplicationSecretStorage,
+		private readonly _accountAccess: IAccountAccess,
 		private readonly _secretStorage: SecretStorage,
 		private readonly _logger: LogOutputChannel,
-		private readonly _cloudName: string
+		disposables: Disposable[]
 	) {
-		this._pcasSecretStorage = new PublicClientApplicationsSecretStorage(_secretStorage, _cloudName);
 		this._disposable = Disposable.from(
-			this._pcasSecretStorage,
+			...disposables,
 			this._registerSecretStorageHandler(),
 			this._onDidAccountsChangeEmitter
 		);
 	}
 
+	static async create(
+		secretStorage: SecretStorage,
+		logger: LogOutputChannel,
+		cloudName: string
+	): Promise<CachedPublicClientApplicationManager> {
+		const pcasSecretStorage = await PublicClientApplicationsSecretStorage.create(secretStorage, cloudName);
+		// TODO: Remove the migrations in a version
+		const migrations = await pcasSecretStorage.getOldValue();
+		const accountAccess = await ScopedAccountAccess.create(secretStorage, cloudName, logger, migrations);
+		const manager = new CachedPublicClientApplicationManager(pcasSecretStorage, accountAccess, secretStorage, logger, [pcasSecretStorage, accountAccess]);
+		await manager.initialize();
+		return manager;
+	}
+
 	private _registerSecretStorageHandler() {
 		return this._pcasSecretStorage.onDidChange(() => this._handleSecretStorageChange());
 	}
 
-	async initialize() {
+	private async initialize() {
 		this._logger.debug('[initialize] Initializing PublicClientApplicationManager');
 		let clientIds: string[] | undefined;
 		try {
-			await this._pcasSecretStorage.initialize();
-			//TODO: Remove this in a version
-			const migrations = await this._pcasSecretStorage.getOldValue();
-			this._accountAccess = new ScopedAccountAccess(this._secretStorage, this._cloudName, this._logger, migrations);
 			clientIds = await this._pcasSecretStorage.get();
 		} catch (e) {
 			// data is corrupted
@@ -124,14 +133,12 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient
 					this._logger.error(`[getOrCreate] [${clientId}] Error migrating refresh token:`, e);
 				}
 			}
-			// reinitialize the PCA so the account is properly cached
-			await pca.initialize();
 		}
 		return pca;
 	}
 
 	private async _doCreatePublicClientApplication(clientId: string): Promise<ICachedPublicClientApplication> {
-		const pca = new CachedPublicClientApplication(clientId, this._secretStorage, this._accountAccess!, this._logger);
+		const pca = await CachedPublicClientApplication.create(clientId, this._secretStorage, this._accountAccess, this._logger);
 		this._pcas.set(clientId, pca);
 		const disposable = Disposable.from(
 			pca,
@@ -147,8 +154,10 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient
 			})
 		);
 		this._pcaDisposables.set(clientId, disposable);
-		// Intialize the PCA after the `onDidAccountsChange` is set so we get initial state.
-		await pca.initialize();
+		// Fire for the initial state and only if accounts exist
+		if (pca.accounts.length > 0) {
+			this._onDidAccountsChangeEmitter.fire({ added: pca.accounts, changed: [], deleted: [] });
+		}
 		return pca;
 	}
 
@@ -205,7 +214,15 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient
 	}
 }
 
-class PublicClientApplicationsSecretStorage {
+interface IPublicClientApplicationSecretStorage {
+	get(): Promise<string[] | undefined>;
+	getOldValue(): Promise<{ clientId: string; authority: string }[] | undefined>;
+	store(value: string[]): Thenable<void>;
+	delete(): Thenable<void>;
+	onDidChange: Event<void>;
+}
+
+class PublicClientApplicationsSecretStorage implements IPublicClientApplicationSecretStorage, Disposable {
 	private _disposable: Disposable;
 
 	private readonly _onDidChangeEmitter = new EventEmitter<void>;
@@ -214,7 +231,7 @@ class PublicClientApplicationsSecretStorage {
 	private readonly _oldKey = `publicClientApplications-${this._cloudName}`;
 	private readonly _key = `publicClients-${this._cloudName}`;
 
-	constructor(private readonly _secretStorage: SecretStorage, private readonly _cloudName: string) {
+	private constructor(private readonly _secretStorage: SecretStorage, private readonly _cloudName: string) {
 		this._disposable = Disposable.from(
 			this._onDidChangeEmitter,
 			this._secretStorage.onDidChange(e => {
@@ -225,11 +242,17 @@ class PublicClientApplicationsSecretStorage {
 		);
 	}
 
+	static async create(secretStorage: SecretStorage, cloudName: string): Promise<PublicClientApplicationsSecretStorage> {
+		const storage = new PublicClientApplicationsSecretStorage(secretStorage, cloudName);
+		await storage.initialize();
+		return storage;
+	}
+
 	/**
 	 * Runs the migration.
 	 * TODO: Remove this after a version.
 	 */
-	async initialize() {
+	private async initialize() {
 		const oldValue = await this.getOldValue();
 		if (!oldValue) {
 			return;