Extension host veto is registered multiple times on restart (fix microsoft#183778) (microsoft#184127)

Extension host veto is registered multiple times on restart (microsoft#183778)

Author: bpasero

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

@@ -106,26 +106,6 @@ 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) {
-						// Veto
-						return true;
-					}
-				}
-				return false; // Don't veto
-			})(), localize('vetoExtHostRestart', "One or more custom editors could not be saved."));
-		}));
 	}
 
 	public $registerTextEditorProvider(extensionData: extHostProtocol.WebviewExtensionDescription, viewType: string, options: extHostProtocol.IWebviewPanelOptions, capabilities: extHostProtocol.CustomTextEditorCapabilities, serializeBuffersForPostMessage: boolean): void {
@@ -399,13 +379,29 @@ class MainThreadCustomEditorModel extends ResourceWorkingCopy implements ICustom
 		@IWorkbenchEnvironmentService private readonly _environmentService: IWorkbenchEnvironmentService,
 		@IWorkingCopyService workingCopyService: IWorkingCopyService,
 		@IPathService private readonly _pathService: IPathService,
+		@IExtensionService extensionService: IExtensionService,
 	) {
 		super(MainThreadCustomEditorModel.toWorkingCopyResource(_viewType, _editorResource), fileService);
 
 		this._fromBackup = fromBackup;
 
 		if (_editable) {
 			this._register(workingCopyService.registerWorkingCopy(this));
+
+			this._register(extensionService.onWillStop(e => {
+				if (!this.isDirty()) {
+					return;
+				}
+
+				e.veto((async () => {
+					const didSave = await this.save();
+					if (!didSave) {
+						// Veto
+						return true;
+					}
+					return false; // Don't veto
+				})(), localize('vetoExtHostRestart', "Custom editor '{0}' could not be saved.", this.name));
+			}));
 		}
 
 		// Normally means we're re-opening an untitled file

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

@@ -70,7 +70,7 @@ export class SimpleNotebookEditorModel extends EditorModel implements INotebookE
 					return true;
 				}
 				return false; // Don't veto
-			})(), localize('vetoExtHostRestart', "Notebook could not be saved."));
+			})(), localize('vetoExtHostRestart', "Notebook '{0}' could not be saved.", this._workingCopy?.name));
 		}));
 	}