Adopt ext host restart for custom editors

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

@@ -20,6 +20,7 @@ import { IFileDialogService } from 'vs/platform/dialogs/common/dialogs';
 import { FileOperation, IFileService } from 'vs/platform/files/common/files';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ILabelService } from 'vs/platform/label/common/label';
+import { INotificationService } from 'vs/platform/notification/common/notification';
 import { IStorageService } from 'vs/platform/storage/common/storage';
 import { IUndoRedoService, UndoRedoElementType } from 'vs/platform/undoRedo/common/undoRedo';
 import { MainThreadWebviewPanels } from 'vs/workbench/api/browser/mainThreadWebviewPanels';
@@ -65,14 +66,15 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 		private readonly mainThreadWebview: MainThreadWebviews,
 		private readonly mainThreadWebviewPanels: MainThreadWebviewPanels,
 		@IExtensionService extensionService: IExtensionService,
+		@INotificationService _notificationService: INotificationService,
 		@IStorageService storageService: IStorageService,
 		@IWorkingCopyService workingCopyService: IWorkingCopyService,
 		@IWorkingCopyFileService workingCopyFileService: IWorkingCopyFileService,
 		@ICustomEditorService private readonly _customEditorService: ICustomEditorService,
 		@IEditorGroupsService private readonly _editorGroupService: IEditorGroupsService,
-		@IWebviewWorkbenchService private readonly _webviewWorkbenchService: IWebviewWorkbenchService,
-		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@IEditorService private readonly _editorService: IEditorService,
+		@IInstantiationService private readonly _instantiationService: IInstantiationService,
+		@IWebviewWorkbenchService private readonly _webviewWorkbenchService: IWebviewWorkbenchService,
 	) {
 		super();
 
@@ -106,6 +108,28 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 
 		// Working copy operations
 		this._register(workingCopyFileService.onWillRunWorkingCopyFileOperation(async e => this.onWillRunWorkingCopyFileOperation(e)));
+
+		this._register(extensionService.onWillStop(e => {
+			const dirtyCustomEditors = workingCopyService.workingCopies.filter(workingCopy => {
+				return workingCopy instanceof MainThreadCustomEditorModel && workingCopy.isDirty();
+			});
+			if (!dirtyCustomEditors.length) {
+				return;
+			}
+
+			e.veto((async () => {
+				for (const dirtyCustomEditor of dirtyCustomEditors) {
+					const didSave = await dirtyCustomEditor.save();
+					if (!didSave) {
+						_notificationService.error(localize('vetoExtHostRestart', "'{0}' was not saved which prevented restarting the extension host. Please save or revert this file to restart the extension host.", dirtyCustomEditor.name));
+
+						// Veto
+						return true;
+					}
+				}
+				return false; // Don't veto
+			})(), 'mainThreadCustomEditors');
+		}));
 	}
 
 	public $registerTextEditorProvider(extensionData: extHostProtocol.WebviewExtensionDescription, viewType: string, options: extHostProtocol.IWebviewPanelOptions, capabilities: extHostProtocol.CustomTextEditorCapabilities, serializeBuffersForPostMessage: boolean): void {