Adopt new extension host restart lifecycle (microsoft#180514) (microsoft#183293)

* Adopt new extension host restart lifecycle (microsoft#180514)

* adopt for notebooks

* fix tests

Author: bpasero

src/vs/workbench/contrib/extensions/browser/extensionEnablementWorkspaceTrustTransitionParticipant.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { localize } from 'vs/nls';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { IWorkspaceTrustEnablementService, IWorkspaceTrustManagementService, IWorkspaceTrustTransitionParticipant } from 'vs/platform/workspace/common/workspaceTrust';
 import { IWorkbenchContribution } from 'vs/workbench/common/contributions';
@@ -38,9 +39,11 @@ export class ExtensionEnablementWorkspaceTrustTransitionParticipant extends Disp
 							if (environmentService.remoteAuthority) {
 								hostService.reload();
 							} else {
-								extensionService.stopExtensionHosts(true); // TODO@lszomoru adopt support for extension host to veto stopping
+								const stopped = await extensionService.stopExtensionHosts(localize('restartExtensionHost.reason', "Restarting extension host due to workspace trust change."));
 								await extensionEnablementService.updateExtensionsEnablementsWhenWorkspaceTrustChanges();
-								extensionService.startExtensionHosts();
+								if (stopped) {
+									extensionService.startExtensionHosts();
+								}
 							}
 						}
 					}

src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts

@@ -12,11 +12,13 @@ import { Schemas } from 'vs/base/common/network';
 import { filter } from 'vs/base/common/objects';
 import { assertType } from 'vs/base/common/types';
 import { URI } from 'vs/base/common/uri';
+import { localize } from 'vs/nls';
 import { IRevertOptions, ISaveOptions, IUntypedEditorInput } from 'vs/workbench/common/editor';
 import { EditorModel } from 'vs/workbench/common/editor/editorModel';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { ICellDto2, INotebookEditorModel, INotebookLoadOptions, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookData } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookSerializer, INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
+import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 import { IFilesConfigurationService } from 'vs/workbench/services/filesConfiguration/common/filesConfigurationService';
 import { ILifecycleService } from 'vs/workbench/services/lifecycle/common/lifecycle';
 import { IFileWorkingCopyManager } from 'vs/workbench/services/workingCopy/common/fileWorkingCopyManager';
@@ -47,13 +49,29 @@ export class SimpleNotebookEditorModel extends EditorModel implements INotebookE
 		readonly viewType: string,
 		private readonly _workingCopyManager: IFileWorkingCopyManager<NotebookFileWorkingCopyModel, NotebookFileWorkingCopyModel>,
 		@ILifecycleService lifecycleService: ILifecycleService,
-		@IFilesConfigurationService private readonly _filesConfigurationService: IFilesConfigurationService
+		@IFilesConfigurationService private readonly _filesConfigurationService: IFilesConfigurationService,
+		@IExtensionService extensionService: IExtensionService
 	) {
 		super();
 
 		if (this.viewType === 'interactive') {
 			lifecycleService.onBeforeShutdown(async e => e.veto(this.onBeforeShutdown(), 'veto.InteractiveWindow'));
 		}
+
+		this._register(extensionService.onWillStop(e => {
+			if (!this.isDirty()) {
+				return;
+			}
+
+			e.veto((async () => {
+				const didSave = await this._workingCopy?.save();
+				if (!didSave) {
+					// Veto
+					return true;
+				}
+				return false; // Don't veto
+			})(), localize('vetoExtHostRestart', "Notebook could not be saved."));
+		}));
 	}
 
 	private async onBeforeShutdown() {

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

@@ -204,7 +204,7 @@ export class WorkspaceChangeExtHostRelauncher extends Disposable implements IWor
 			if (environmentService.remoteAuthority) {
 				hostService.reload(); // TODO@aeschli, workaround
 			} else if (isNative) {
-				const stopped = await extensionService.stopExtensionHosts(localize('restartExtensionHost.reason', "Restart Extension Host."));
+				const stopped = await extensionService.stopExtensionHosts(localize('restartExtensionHost.reason', "Restarting extension host due to a workspace folder change."));
 				if (stopped) {
 					extensionService.startExtensionHosts();
 				}

src/vs/workbench/services/extensions/common/abstractExtensionService.ts

@@ -623,14 +623,8 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 
 	//#region Stopping / Starting / Restarting
 
-	public stopExtensionHosts(reason: string): Promise<boolean>;
-	public stopExtensionHosts(force: true): void;
-	public stopExtensionHosts(arg0: true | string): void | Promise<boolean> {
-		if (arg0 === true) {
-			return this._doStopExtensionHosts();
-		}
-
-		return this._doStopExtensionHostsWithVeto(arg0);
+	public stopExtensionHosts(reason: string): Promise<boolean> {
+		return this._doStopExtensionHostsWithVeto(reason);
 	}
 
 	protected _doStopExtensionHosts(): void {

src/vs/workbench/services/extensions/common/extensions.ts

@@ -460,11 +460,6 @@ export interface IExtensionService {
 	 */
 	stopExtensionHosts(reason: string): Promise<boolean>;
 
-	/**
-	 * @deprecated Use `stopExtensionHosts()` instead.
-	 */
-	stopExtensionHosts(force: true): void;
-
 	/**
 	 * Starts the extension hosts.
 	 */

src/vs/workbench/services/extensions/electron-sandbox/nativeExtensionService.ts

@@ -723,7 +723,7 @@ class RestartExtensionHostAction extends Action2 {
 	async run(accessor: ServicesAccessor): Promise<void> {
 		const extensionService = accessor.get(IExtensionService);
 
-		const stopped = await extensionService.stopExtensionHosts(nls.localize('restartExtensionHost.reason', "Restart Extension Host."));
+		const stopped = await extensionService.stopExtensionHosts(nls.localize('restartExtensionHost.reason', "Restarting extension host on explicit request."));
 		if (stopped) {
 			extensionService.startExtensionHosts();
 		}

src/vs/workbench/services/extensions/test/browser/extensionService.test.ts

@@ -256,13 +256,13 @@ suite('ExtensionService', () => {
 
 	test('issue #152204: Remote extension host not disposed after closing vscode client', async () => {
 		await extService.startExtensionHosts();
-		extService.stopExtensionHosts(true);
+		await extService.stopExtensionHosts('foo');
 		assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3', 'dispose 3', 'dispose 2', 'dispose 1']));
 	});
 
 	test('Extension host disposed when awaited', async () => {
 		await extService.startExtensionHosts();
-		await extService.stopExtensionHosts(`foo`);
+		await extService.stopExtensionHosts('foo');
 		assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3', 'dispose 3', 'dispose 2', 'dispose 1']));
 	});
 
@@ -272,7 +272,7 @@ suite('ExtensionService', () => {
 		extService.onWillStop(e => e.veto(true, 'test 1'));
 		extService.onWillStop(e => e.veto(false, 'test 2'));
 
-		await extService.stopExtensionHosts(`foo`);
+		await extService.stopExtensionHosts('foo');
 		assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3']));
 	});
 
@@ -283,7 +283,7 @@ suite('ExtensionService', () => {
 		extService.onWillStop(e => e.veto(Promise.resolve(true), 'test 2'));
 		extService.onWillStop(e => e.veto(Promise.resolve(false), 'test 3'));
 
-		await extService.stopExtensionHosts(`foo`);
+		await extService.stopExtensionHosts('foo');
 		assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3']));
 	});
 });

src/vs/workbench/services/userDataProfile/browser/userDataProfileManagement.ts

@@ -118,7 +118,7 @@ export class UserDataProfileManagementService extends Disposable implements IUse
 		const isRemoteWindow = !!this.environmentService.remoteAuthority;
 
 		if (!isRemoteWindow) {
-			if (!(await this.extensionService.stopExtensionHosts(localize('switch profile', "Switching Profile")))) {
+			if (!(await this.extensionService.stopExtensionHosts(localize('switch profile', "Switching to a profile.")))) {
 				throw new CancellationError();
 			}
 		}