Ask for edit sessions auth when coming from Continue On (microsoft#160843)

* Ask for edit sessions auth when coming from Continue On

* Only append `continueOn` query param if there is an edit session

* Persist continueOn for new windows

* Don't double prompt to sign in

* Add badge if user didn't sign in

* Add logging and dev commands

* Log more info

* More logging

* 🤔

* Cleanup

* Actually check for pending edit sessions

* More cleanup

* Update test

Author: joyceerhl

src/vs/code/electron-main/app.ts

@@ -918,6 +918,14 @@ export class CodeApplication extends Disposable {
 				// or if no window is open (macOS only)
 				shouldOpenInNewWindow ||= isMacintosh && windowsMainService.getWindowCount() === 0;
 
+				// Pass along whether the application is being opened via a Continue On flow
+				const continueOn = params.get('continueOn');
+				if (continueOn !== null) {
+					environmentService.continueOn = continueOn ?? undefined;
+					params.delete('continueOn');
+					uri = uri.with({ query: params.toString() });
+				}
+
 				// Check for URIs to open in window
 				const windowOpenableFromProtocolLink = app.getWindowOpenableFromProtocolLink(uri);
 				logService.trace('app#handleURL: windowOpenableFromProtocolLink = ', windowOpenableFromProtocolLink);

src/vs/platform/environment/common/argv.ts

@@ -90,6 +90,7 @@ export interface NativeParsedArgs {
 	'logsPath'?: string;
 	'__enable-file-policy'?: boolean;
 	editSessionId?: string;
+	continueOn?: string;
 	'locate-shell-integration-path'?: string;
 	'profile'?: string;
 	'profile-temp'?: boolean;

src/vs/platform/environment/common/environment.ts

@@ -65,6 +65,7 @@ export interface IEnvironmentService {
 	sync: 'on' | 'off' | undefined;
 
 	// --- continue edit session
+	continueOn?: string;
 	editSessionId?: string;
 	editSessionsLogResource: URI;
 

src/vs/platform/environment/common/environmentService.ts

@@ -247,6 +247,14 @@ export abstract class AbstractNativeEnvironmentService implements INativeEnviron
 
 	editSessionId: string | undefined = this.args['editSessionId'];
 
+	get continueOn(): string | undefined {
+		return this.args['continueOn'];
+	}
+
+	set continueOn(value: string | undefined) {
+		this.args['continueOn'] = value;
+	}
+
 	get args(): NativeParsedArgs { return this._args; }
 
 	constructor(

src/vs/platform/environment/node/argv.ts

@@ -129,6 +129,7 @@ export const OPTIONS: OptionDescriptions<Required<NativeParsedArgs>> = {
 	'logsPath': { type: 'string' },
 	'__enable-file-policy': { type: 'boolean' },
 	'editSessionId': { type: 'string' },
+	'continueOn': { type: 'string' },
 	'locate-shell-integration-path': { type: 'string', args: ['bash', 'pwsh', 'zsh', 'fish'] },
 
 	'enable-coi': { type: 'boolean' },

src/vs/platform/windows/electron-main/windowImpl.ts

@@ -973,6 +973,7 @@ export class CodeWindow extends Disposable implements ICodeWindow {
 
 		configuration.isInitialStartup = false; // since this is a reload
 		configuration.policiesData = this.policyService.serialize(); // set policies data again
+		configuration.continueOn = this.environmentMainService.continueOn;
 		configuration.profiles = {
 			all: this.userDataProfilesService.profiles,
 			profile: this.profile || this.userDataProfilesService.defaultProfile

src/vs/platform/windows/electron-main/windowsMainService.ts

@@ -1367,6 +1367,7 @@ export class WindowsMainService extends Disposable implements IWindowsMainServic
 			accessibilitySupport: app.accessibilitySupportEnabled,
 			colorScheme: this.themeMainService.getColorScheme(),
 			policiesData: this.policyService.serialize(),
+			continueOn: this.environmentMainService.continueOn,
 		};
 
 		let window: ICodeWindow | undefined;

src/vs/workbench/contrib/editSessions/browser/editSessions.contribution.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, MutableDisposable } from 'vs/base/common/lifecycle';
 import { IWorkbenchContributionsRegistry, Extensions as WorkbenchExtensions, IWorkbenchContribution } from 'vs/workbench/common/contributions';
 import { Registry } from 'vs/platform/registry/common/platform';
 import { ILifecycleService, LifecyclePhase } from 'vs/workbench/services/lifecycle/common/lifecycle';
@@ -54,6 +54,8 @@ import { ThemeIcon } from 'vs/platform/theme/common/themeService';
 import { IOutputService } from 'vs/workbench/services/output/common/output';
 import * as Constants from 'vs/workbench/contrib/logs/common/logConstants';
 import { sha1Hex } from 'vs/base/browser/hash';
+import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
+import { IActivityService, NumberBadge } from 'vs/workbench/services/activity/common/activity';
 
 registerSingleton(IEditSessionsLogService, EditSessionsLogService, false);
 registerSingleton(IEditSessionsStorageService, EditSessionsWorkbenchService, false);
@@ -89,6 +91,9 @@ export class EditSessionsContribution extends Disposable implements IWorkbenchCo
 
 	private readonly shouldShowViewsContext: IContextKey<boolean>;
 
+	private static APPLICATION_LAUNCHED_VIA_CONTINUE_ON_STORAGE_KEY = 'applicationLaunchedViaContinueOn';
+	private accountsMenuBadgeDisposable = this._register(new MutableDisposable());
+
 	constructor(
 		@IEditSessionsStorageService private readonly editSessionsStorageService: IEditSessionsStorageService,
 		@IFileService private readonly fileService: IFileService,
@@ -109,7 +114,9 @@ export class EditSessionsContribution extends Disposable implements IWorkbenchCo
 		@ICommandService private commandService: ICommandService,
 		@IContextKeyService private readonly contextKeyService: IContextKeyService,
 		@IFileDialogService private readonly fileDialogService: IFileDialogService,
-		@ILifecycleService private readonly lifecycleService: ILifecycleService
+		@ILifecycleService private readonly lifecycleService: ILifecycleService,
+		@IStorageService private readonly storageService: IStorageService,
+		@IActivityService private readonly activityService: IActivityService,
 	) {
 		super();
 
@@ -123,6 +130,7 @@ export class EditSessionsContribution extends Disposable implements IWorkbenchCo
 
 		this._register(this.fileService.registerProvider(EditSessionsFileSystemProvider.SCHEMA, new EditSessionsFileSystemProvider(this.editSessionsStorageService)));
 		this.lifecycleService.onWillShutdown((e) => e.join(this.autoStoreEditSession(), { id: 'autoStoreEditSession', label: localize('autoStoreEditSession', 'Storing current edit session...') }));
+		this._register(this.editSessionsStorageService.onDidSignIn(() => this.updateAccountsMenuBadge()));
 	}
 
 	private autoResumeEditSession() {
@@ -136,21 +144,58 @@ export class EditSessionsContribution extends Disposable implements IWorkbenchCo
 			this.telemetryService.publicLog2<ResumeEvent, ResumeClassification>('editSessions.continue.resume');
 
 			if (this.environmentService.editSessionId !== undefined) {
+				this.logService.info(`Resuming edit session, reason: found editSessionId ${this.environmentService.editSessionId} in environment service...`);
 				await this.resumeEditSession(this.environmentService.editSessionId).finally(() => this.environmentService.editSessionId = undefined);
 			} else if (
 				this.configurationService.getValue('workbench.editSessions.autoResume') === 'onReload' &&
 				this.editSessionsStorageService.isSignedIn
 			) {
+				this.logService.info('Resuming edit session, reason: edit sessions enabled...');
 				// Attempt to resume edit session based on edit workspace identifier
 				// Note: at this point if the user is not signed into edit sessions,
 				// we don't want them to be prompted to sign in and should just return early
 				await this.resumeEditSession(undefined, true);
+			} else {
+				// The application has previously launched via a protocol URL Continue On flow
+				const hasApplicationLaunchedFromContinueOnFlow = this.storageService.getBoolean(EditSessionsContribution.APPLICATION_LAUNCHED_VIA_CONTINUE_ON_STORAGE_KEY, StorageScope.APPLICATION, false);
+
+				if ((this.environmentService.continueOn !== undefined) &&
+					!this.editSessionsStorageService.isSignedIn &&
+					// and user has not yet been prompted to sign in on this machine
+					hasApplicationLaunchedFromContinueOnFlow === false
+				) {
+					this.storageService.store(EditSessionsContribution.APPLICATION_LAUNCHED_VIA_CONTINUE_ON_STORAGE_KEY, true, StorageScope.APPLICATION, StorageTarget.MACHINE);
+					await this.editSessionsStorageService.initialize(true);
+					if (this.editSessionsStorageService.isSignedIn) {
+						await this.resumeEditSession(undefined, true);
+					} else {
+						this.updateAccountsMenuBadge();
+					}
+					// store the fact that we prompted the user
+				} else if (!this.editSessionsStorageService.isSignedIn &&
+					// and user has been prompted to sign in on this machine
+					hasApplicationLaunchedFromContinueOnFlow === true
+				) {
+					// display a badge in the accounts menu but do not prompt the user to sign in again
+					this.updateAccountsMenuBadge();
+					// attempt a resume if we are in a pending state and the user just signed in
+					this._register(this.editSessionsStorageService.onDidSignIn(async () => this.resumeEditSession(undefined, true)));
+				}
 			}
 
 			performance.mark('code/didResumeEditSessionFromIdentifier');
 		});
 	}
 
+	private updateAccountsMenuBadge() {
+		if (this.editSessionsStorageService.isSignedIn) {
+			return this.accountsMenuBadgeDisposable.clear();
+		}
+
+		const badge = new NumberBadge(1, () => localize('check for pending edit sessions', 'Check for pending edit sessions'));
+		this.accountsMenuBadgeDisposable.value = this.activityService.showAccountsActivity({ badge, priority: 1 });
+	}
+
 	private async autoStoreEditSession() {
 		if (this.configurationService.getValue('workbench.experimental.editSessions.autoStore') === 'onShutdown') {
 			await this.progressService.withProgress({
@@ -252,7 +297,7 @@ export class EditSessionsContribution extends Disposable implements IWorkbenchCo
 				if (ref !== undefined && uri !== 'noDestinationUri') {
 					const encodedRef = encodeURIComponent(ref);
 					uri = uri.with({
-						query: uri.query.length > 0 ? (uri.query + `&${queryParamName}=${encodedRef}`) : `${queryParamName}=${encodedRef}`
+						query: uri.query.length > 0 ? (uri.query + `&${queryParamName}=${encodedRef}&continueOn=1`) : `${queryParamName}=${encodedRef}&continueOn=1`
 					});
 
 					// Open the URI

src/vs/workbench/contrib/editSessions/browser/editSessionsStorageService.ts

@@ -27,6 +27,7 @@ import { isWeb } from 'vs/base/common/platform';
 import { ICommandService } from 'vs/platform/commands/common/commands';
 import { Codicon } from 'vs/base/common/codicons';
 import { IUserDataSyncMachinesService, UserDataSyncMachinesService } from 'vs/platform/userDataSync/common/userDataSyncMachines';
+import { Emitter } from 'vs/base/common/event';
 
 type ExistingSession = IQuickPickItem & { session: AuthenticationSession & { providerId: string } };
 type AuthenticationProviderOption = IQuickPickItem & { provider: IAuthenticationProvider };
@@ -50,6 +51,11 @@ export class EditSessionsWorkbenchService extends Disposable implements IEditSes
 		return this.existingSessionId !== undefined;
 	}
 
+	private _didSignIn = new Emitter<void>();
+	get onDidSignIn() {
+		return this._didSignIn.event;
+	}
+
 	constructor(
 		@IFileService private readonly fileService: IFileService,
 		@IStorageService private readonly storageService: IStorageService,
@@ -162,6 +168,9 @@ export class EditSessionsWorkbenchService extends Disposable implements IEditSes
 		}
 		this.initialized = await this.doInitialize(fromContinueOn);
 		this.signedInContext.set(this.initialized);
+		if (this.initialized) {
+			this._didSignIn.fire();
+		}
 		return this.initialized;
 
 	}

src/vs/workbench/contrib/editSessions/common/editSessions.ts

@@ -12,6 +12,7 @@ import { createDecorator } from 'vs/platform/instantiation/common/instantiation'
 import { ILogService } from 'vs/platform/log/common/log';
 import { registerIcon } from 'vs/platform/theme/common/iconRegistry';
 import { IResourceRefHandle } from 'vs/platform/userDataSync/common/userDataSync';
+import { Event } from 'vs/base/common/event';
 
 export const EDIT_SESSION_SYNC_CATEGORY: ILocalizedString = {
 	original: 'Edit Sessions',
@@ -23,6 +24,7 @@ export interface IEditSessionsStorageService {
 	_serviceBrand: undefined;
 
 	readonly isSignedIn: boolean;
+	readonly onDidSignIn: Event<void>;
 
 	initialize(fromContinueOn: boolean): Promise<boolean>;
 	read(ref: string | undefined): Promise<{ ref: string; editSession: EditSession } | undefined>;