Implement getAccounts API over getSessions (#215874)

And plumb that through to the Microsoft auth provider

extensions/microsoft-authentication/package.json

@@ -14,7 +14,8 @@
   ],
   "activationEvents": [],
   "enabledApiProposals": [
-    "idToken"
+    "idToken",
+    "authGetSessions"
   ],
   "capabilities": {
     "virtualWorkspaces": true,

extensions/microsoft-authentication/src/AADHelper.ts

@@ -203,11 +203,13 @@ export class AzureActiveDirectoryService {
 		return this._sessionChangeEmitter.event;
 	}
 
-	public getSessions(scopes?: string[]): Promise<vscode.AuthenticationSession[]> {
+	public getSessions(scopes?: string[], account?: vscode.AuthenticationSessionAccountInformation): Promise<vscode.AuthenticationSession[]> {
 		if (!scopes) {
 			this._logger.info('Getting sessions for all scopes...');
-			const sessions = this._tokens.map(token => this.convertToSessionSync(token));
-			this._logger.info(`Got ${sessions.length} sessions for all scopes...`);
+			const sessions = this._tokens
+				.filter(token => !account?.label || token.account.label === account.label)
+				.map(token => this.convertToSessionSync(token));
+			this._logger.info(`Got ${sessions.length} sessions for all scopes${account ? ` for account '${account.label}'` : ''}...`);
 			return Promise.resolve(sessions);
 		}
 
@@ -238,23 +240,43 @@ export class AzureActiveDirectoryService {
 			tenant: this.getTenantId(scopes),
 		};
 
-		this._logger.trace(`[${scopeData.scopeStr}] Queued getting sessions`);
-		return this._sequencer.queue(modifiedScopesStr, () => this.doGetSessions(scopeData));
+		this._logger.trace(`[${scopeData.scopeStr}] Queued getting sessions` + account ? ` for ${account?.label}` : '');
+		return this._sequencer.queue(modifiedScopesStr, () => this.doGetSessions(scopeData, account));
 	}
 
-	private async doGetSessions(scopeData: IScopeData): Promise<vscode.AuthenticationSession[]> {
-		this._logger.info(`[${scopeData.scopeStr}] Getting sessions`);
+	private async doGetSessions(scopeData: IScopeData, account?: vscode.AuthenticationSessionAccountInformation): Promise<vscode.AuthenticationSession[]> {
+		this._logger.info(`[${scopeData.scopeStr}] Getting sessions` + account ? ` for ${account?.label}` : '');
 
-		const matchingTokens = this._tokens.filter(token => token.scope === scopeData.scopeStr);
+		const matchingTokens = this._tokens
+			.filter(token => token.scope === scopeData.scopeStr)
+			.filter(token => !account?.label || token.account.label === account.label);
 		// If we still don't have a matching token try to get a new token from an existing token by using
 		// the refreshToken. This is documented here:
 		// https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-auth-code-flow#refresh-the-access-token
 		// "Refresh tokens are valid for all permissions that your client has already received consent for."
 		if (!matchingTokens.length) {
-			// Get a token with the correct client id.
-			const token = scopeData.clientId === DEFAULT_CLIENT_ID
-				? this._tokens.find(t => t.refreshToken && !t.scope.includes('VSCODE_CLIENT_ID'))
-				: this._tokens.find(t => t.refreshToken && t.scope.includes(`VSCODE_CLIENT_ID:${scopeData.clientId}`));
+			// Get a token with the correct client id and account.
+			let token: IToken | undefined;
+			for (const t of this._tokens) {
+				// No refresh token, so we can't make a new token from this session
+				if (!t.refreshToken) {
+					continue;
+				}
+				// Need to make sure the account matches if we were provided one
+				if (account?.label && t.account.label !== account.label) {
+					continue;
+				}
+				// If the client id is the default client id, then check for the absence of the VSCODE_CLIENT_ID scope
+				if (scopeData.clientId === DEFAULT_CLIENT_ID && !t.scope.includes('VSCODE_CLIENT_ID')) {
+					token = t;
+					break;
+				}
+				// If the client id is not the default client id, then check for the matching VSCODE_CLIENT_ID scope
+				if (scopeData.clientId !== DEFAULT_CLIENT_ID && t.scope.includes(`VSCODE_CLIENT_ID:${scopeData.clientId}`)) {
+					token = t;
+					break;
+				}
+			}
 
 			if (token) {
 				this._logger.trace(`[${scopeData.scopeStr}] '${token.sessionId}' Found a matching token with a different scopes '${token.scope}'. Attempting to get a new session using the existing session.`);
@@ -275,7 +297,7 @@ export class AzureActiveDirectoryService {
 			.map(result => (result as PromiseFulfilledResult<vscode.AuthenticationSession>).value);
 	}
 
-	public createSession(scopes: string[]): Promise<vscode.AuthenticationSession> {
+	public createSession(scopes: string[], account?: vscode.AuthenticationSessionAccountInformation): Promise<vscode.AuthenticationSession> {
 		let modifiedScopes = [...scopes];
 		if (!modifiedScopes.includes('openid')) {
 			modifiedScopes.push('openid');
@@ -301,11 +323,11 @@ export class AzureActiveDirectoryService {
 		};
 
 		this._logger.trace(`[${scopeData.scopeStr}] Queued creating session`);
-		return this._sequencer.queue(scopeData.scopeStr, () => this.doCreateSession(scopeData));
+		return this._sequencer.queue(scopeData.scopeStr, () => this.doCreateSession(scopeData, account));
 	}
 
-	private async doCreateSession(scopeData: IScopeData): Promise<vscode.AuthenticationSession> {
-		this._logger.info(`[${scopeData.scopeStr}] Creating session`);
+	private async doCreateSession(scopeData: IScopeData, account?: vscode.AuthenticationSessionAccountInformation): Promise<vscode.AuthenticationSession> {
+		this._logger.info(`[${scopeData.scopeStr}] Creating session` + account ? ` for ${account?.label}` : '');
 
 		const runsRemote = vscode.env.remoteName !== undefined;
 		const runsServerless = vscode.env.remoteName === undefined && vscode.env.uiKind === vscode.UIKind.Web;
@@ -316,25 +338,25 @@ export class AzureActiveDirectoryService {
 
 		return await vscode.window.withProgress({ location: vscode.ProgressLocation.Notification, title: vscode.l10n.t('Signing in to your account...'), cancellable: true }, async (_progress, token) => {
 			if (runsRemote || runsServerless) {
-				return await this.createSessionWithoutLocalServer(scopeData, token);
+				return await this.createSessionWithoutLocalServer(scopeData, account?.label, token);
 			}
 
 			try {
-				return await this.createSessionWithLocalServer(scopeData, token);
+				return await this.createSessionWithLocalServer(scopeData, account?.label, token);
 			} catch (e) {
 				this._logger.error(`[${scopeData.scopeStr}] Error creating session: ${e}`);
 
 				// If the error was about starting the server, try directly hitting the login endpoint instead
 				if (e.message === 'Error listening to server' || e.message === 'Closed' || e.message === 'Timeout waiting for port') {
-					return this.createSessionWithoutLocalServer(scopeData, token);
+					return this.createSessionWithoutLocalServer(scopeData, account?.label, token);
 				}
 
 				throw e;
 			}
 		});
 	}
 
-	private async createSessionWithLocalServer(scopeData: IScopeData, token: vscode.CancellationToken): Promise<vscode.AuthenticationSession> {
+	private async createSessionWithLocalServer(scopeData: IScopeData, loginHint: string | undefined, token: vscode.CancellationToken): Promise<vscode.AuthenticationSession> {
 		this._logger.trace(`[${scopeData.scopeStr}] Starting login flow with local server`);
 		const codeVerifier = generateCodeVerifier();
 		const codeChallenge = await generateCodeChallenge(codeVerifier);
@@ -344,11 +366,15 @@ export class AzureActiveDirectoryService {
 			client_id: scopeData.clientId,
 			redirect_uri: redirectUrl,
 			scope: scopeData.scopesToSend,
-			prompt: 'select_account',
 			code_challenge_method: 'S256',
 			code_challenge: codeChallenge,
-		}).toString();
-		const loginUrl = new URL(`${scopeData.tenant}/oauth2/v2.0/authorize?${qs}`, this._env.activeDirectoryEndpointUrl).toString();
+		});
+		if (loginHint) {
+			qs.set('login_hint', loginHint);
+		} else {
+			qs.set('prompt', 'select_account');
+		}
+		const loginUrl = new URL(`${scopeData.tenant}/oauth2/v2.0/authorize?${qs.toString()}`, this._env.activeDirectoryEndpointUrl).toString();
 		const server = new LoopbackAuthServer(path.join(__dirname, '../media'), loginUrl);
 		await server.start();
 
@@ -370,7 +396,7 @@ export class AzureActiveDirectoryService {
 		return session;
 	}
 
-	private async createSessionWithoutLocalServer(scopeData: IScopeData, token: vscode.CancellationToken): Promise<vscode.AuthenticationSession> {
+	private async createSessionWithoutLocalServer(scopeData: IScopeData, loginHint: string | undefined, token: vscode.CancellationToken): Promise<vscode.AuthenticationSession> {
 		this._logger.trace(`[${scopeData.scopeStr}] Starting login flow without local server`);
 		let callbackUri = await vscode.env.asExternalUri(vscode.Uri.parse(`${vscode.env.uriScheme}://vscode.microsoft-authentication`));
 		const nonce = generateCodeVerifier();
@@ -383,17 +409,22 @@ export class AzureActiveDirectoryService {
 		const codeVerifier = generateCodeVerifier();
 		const codeChallenge = await generateCodeChallenge(codeVerifier);
 		const signInUrl = new URL(`${scopeData.tenant}/oauth2/v2.0/authorize`, this._env.activeDirectoryEndpointUrl);
-		signInUrl.search = new URLSearchParams({
+		const qs = new URLSearchParams({
 			response_type: 'code',
 			client_id: encodeURIComponent(scopeData.clientId),
 			response_mode: 'query',
 			redirect_uri: redirectUrl,
 			state,
 			scope: scopeData.scopesToSend,
-			prompt: 'select_account',
 			code_challenge_method: 'S256',
 			code_challenge: codeChallenge,
-		}).toString();
+		});
+		if (loginHint) {
+			qs.append('login_hint', loginHint);
+		} else {
+			qs.append('prompt', 'select_account');
+		}
+		signInUrl.search = qs.toString();
 		const uri = vscode.Uri.parse(signInUrl.toString());
 		vscode.env.openExternal(uri);
 

extensions/microsoft-authentication/src/extension.ts

@@ -123,8 +123,8 @@ export async function activate(context: vscode.ExtensionContext) {
 
 	context.subscriptions.push(vscode.authentication.registerAuthenticationProvider('microsoft', 'Microsoft', {
 		onDidChangeSessions: loginService.onDidChangeSessions,
-		getSessions: (scopes: string[]) => loginService.getSessions(scopes),
-		createSession: async (scopes: string[]) => {
+		getSessions: (scopes: string[], options?: vscode.AuthenticationProviderSessionOptions) => loginService.getSessions(scopes, options?.account),
+		createSession: async (scopes: string[], options?: vscode.AuthenticationProviderSessionOptions) => {
 			try {
 				/* __GDPR__
 					"login" : {
@@ -138,7 +138,7 @@ export async function activate(context: vscode.ExtensionContext) {
 					scopes: JSON.stringify(scopes.map(s => s.replace(/[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}/i, '{guid}'))),
 				});
 
-				return await loginService.createSession(scopes);
+				return await loginService.createSession(scopes, options?.account);
 			} catch (e) {
 				/* __GDPR__
 					"loginFailed" : { "owner": "TylerLeonhardt", "comment": "Used to determine how often users run into issues with the login flow." }

extensions/microsoft-authentication/tsconfig.json

@@ -22,6 +22,7 @@
 	"include": [
 		"src/**/*",
 		"../../src/vscode-dts/vscode.d.ts",
-		"../../src/vscode-dts/vscode.proposed.idToken.d.ts"
+		"../../src/vscode-dts/vscode.proposed.idToken.d.ts",
+		"../../src/vscode-dts/vscode.proposed.authGetSessions.d.ts"
 	]
 }

src/vs/workbench/api/browser/mainThreadAuthentication.ts

@@ -6,7 +6,7 @@
 import { Disposable, DisposableMap } from 'vs/base/common/lifecycle';
 import * as nls from 'vs/nls';
 import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
-import { IAuthenticationCreateSessionOptions, AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationService, IAuthenticationExtensionsService, INTERNAL_AUTH_PROVIDER_PREFIX as INTERNAL_MODEL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication';
+import { IAuthenticationCreateSessionOptions, AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationService, IAuthenticationExtensionsService, INTERNAL_AUTH_PROVIDER_PREFIX as INTERNAL_MODEL_AUTH_PROVIDER_PREFIX, AuthenticationSessionAccount, IAuthenticationProviderSessionOptions } from 'vs/workbench/services/authentication/common/authentication';
 import { ExtHostAuthenticationShape, ExtHostContext, MainContext, MainThreadAuthenticationShape } from '../common/extHost.protocol';
 import { IDialogService, IPromptButton } from 'vs/platform/dialogs/common/dialogs';
 import Severity from 'vs/base/common/severity';
@@ -31,6 +31,7 @@ interface AuthenticationGetSessionOptions {
 	createIfNone?: boolean;
 	forceNewSession?: boolean | AuthenticationForceNewSessionOptions;
 	silent?: boolean;
+	account?: AuthenticationSessionAccount;
 }
 
 export class MainThreadAuthenticationProvider extends Disposable implements IAuthenticationProvider {
@@ -49,8 +50,8 @@ export class MainThreadAuthenticationProvider extends Disposable implements IAut
 		this.onDidChangeSessions = onDidChangeSessionsEmitter.event;
 	}
 
-	async getSessions(scopes?: string[]) {
-		return this._proxy.$getSessions(this.id, scopes);
+	async getSessions(scopes: string[] | undefined, options: IAuthenticationProviderSessionOptions) {
+		return this._proxy.$getSessions(this.id, scopes, options);
 	}
 
 	createSession(scopes: string[], options: IAuthenticationCreateSessionOptions): Promise<AuthenticationSession> {
@@ -159,7 +160,7 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu
 	}
 
 	private async doGetSession(providerId: string, scopes: string[], extensionId: string, extensionName: string, options: AuthenticationGetSessionOptions): Promise<AuthenticationSession | undefined> {
-		const sessions = await this.authenticationService.getSessions(providerId, scopes, true);
+		const sessions = await this.authenticationService.getSessions(providerId, scopes, options.account, true);
 		const provider = this.authenticationService.getProvider(providerId);
 
 		// Error cases
@@ -213,18 +214,16 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu
 
 			let session;
 			if (sessions?.length && !options.forceNewSession) {
-				session = provider.supportsMultipleAccounts
+				session = provider.supportsMultipleAccounts && !options.account
 					? await this.authenticationExtensionsService.selectSession(providerId, extensionId, extensionName, scopes, sessions)
 					: sessions[0];
 			} else {
-				let sessionToRecreate: AuthenticationSession | undefined;
-				if (typeof options.forceNewSession === 'object' && options.forceNewSession.sessionToRecreate) {
-					sessionToRecreate = options.forceNewSession.sessionToRecreate as AuthenticationSession;
-				} else {
+				let account: AuthenticationSessionAccount | undefined = options.account;
+				if (!account) {
 					const sessionIdToRecreate = this.authenticationExtensionsService.getSessionPreference(providerId, extensionId, scopes);
-					sessionToRecreate = sessionIdToRecreate ? sessions.find(session => session.id === sessionIdToRecreate) : undefined;
+					account = sessionIdToRecreate ? sessions.find(session => session.id === sessionIdToRecreate)?.account : undefined;
 				}
-				session = await this.authenticationService.createSession(providerId, scopes, { activateImmediate: true, sessionToRecreate });
+				session = await this.authenticationService.createSession(providerId, scopes, { activateImmediate: true, account });
 			}
 
 			this.authenticationAccessService.updateAllowedExtensions(providerId, session.account.label, [{ id: extensionId, name: extensionName, allowed: true }]);
@@ -262,7 +261,7 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu
 	}
 
 	async $getSessions(providerId: string, scopes: readonly string[], extensionId: string, extensionName: string): Promise<AuthenticationSession[]> {
-		const sessions = await this.authenticationService.getSessions(providerId, [...scopes], true);
+		const sessions = await this.authenticationService.getSessions(providerId, [...scopes], undefined, true);
 		const accessibleSessions = sessions.filter(s => this.authenticationAccessService.isAccessAllowed(providerId, s.account.label, extensionId));
 		if (accessibleSessions.length) {
 			this.sendProviderUsageTelemetry(extensionId, providerId);
@@ -273,6 +272,19 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu
 		return accessibleSessions;
 	}
 
+	async $getAccounts(providerId: string): Promise<ReadonlyArray<AuthenticationSessionAccount>> {
+		const sessions = await this.authenticationService.getSessions(providerId);
+		const accounts = new Array<AuthenticationSessionAccount>();
+		const seenAccounts = new Set<string>();
+		for (const session of sessions) {
+			if (!seenAccounts.has(session.account.label)) {
+				seenAccounts.add(session.account.label);
+				accounts.push(session.account);
+			}
+		}
+		return accounts;
+	}
+
 	private sendProviderUsageTelemetry(extensionId: string, providerId: string): void {
 		type AuthProviderUsageClassification = {
 			owner: 'TylerLeonhardt';

src/vs/workbench/api/browser/mainThreadLanguageModels.ts

@@ -14,7 +14,7 @@ import { ExtHostLanguageModelsShape, ExtHostContext, MainContext, MainThreadLang
 import { ILanguageModelStatsService } from 'vs/workbench/contrib/chat/common/languageModelStats';
 import { ILanguageModelChatMetadata, IChatResponseFragment, ILanguageModelsService, IChatMessage, ILanguageModelChatSelector } from 'vs/workbench/contrib/chat/common/languageModels';
 import { IAuthenticationAccessService } from 'vs/workbench/services/authentication/browser/authenticationAccessService';
-import { AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationProviderCreateSessionOptions, IAuthenticationService, INTERNAL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication';
+import { AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationService, INTERNAL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication';
 import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 
@@ -161,9 +161,9 @@ class LanguageModelAccessAuthProvider implements IAuthenticationProvider {
 		if (this._session) {
 			return [this._session];
 		}
-		return [await this.createSession(scopes || [], {})];
+		return [await this.createSession(scopes || [])];
 	}
-	async createSession(scopes: string[], options: IAuthenticationProviderCreateSessionOptions): Promise<AuthenticationSession> {
+	async createSession(scopes: string[]): Promise<AuthenticationSession> {
 		this._session = this._createFakeSession(scopes);
 		this._onDidChangeSessions.fire({ added: [this._session], changed: [], removed: [] });
 		return this._session;

src/vs/workbench/api/common/extHost.api.impl.ts

@@ -287,12 +287,19 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 				if (typeof options?.forceNewSession === 'object' && options.forceNewSession.learnMore) {
 					checkProposedApiEnabled(extension, 'authLearnMore');
 				}
+				if (options?.account) {
+					checkProposedApiEnabled(extension, 'authGetSessions');
+				}
 				return extHostAuthentication.getSession(extension, providerId, scopes, options as any);
 			},
 			getSessions(providerId: string, scopes: readonly string[]) {
 				checkProposedApiEnabled(extension, 'authGetSessions');
 				return extHostAuthentication.getSessions(extension, providerId, scopes);
 			},
+			getAccounts(providerId: string) {
+				checkProposedApiEnabled(extension, 'authGetSessions');
+				return extHostAuthentication.getAccounts(providerId);
+			},
 			// TODO: remove this after GHPR and Codespaces move off of it
 			async hasSession(providerId: string, scopes: readonly string[]) {
 				checkProposedApiEnabled(extension, 'authSession');