Hook custom editors up to backup restorer (#92629)

#77131

Adds custom editor restoring to the backup restorer. This PR includes:

- Adds a `vscode-custom-editor` scheme that we use internally for custom editors. This ensures that each custom editor has it's own editing history (even if another custom editor opens the same resource).

- Make the backup restorer understand how to restore  `vscode-custom-editor` editors. This is done by adding a `IEditorInputFactoryRegistry` that lets the custom editor hook into the backup restorer

src/vs/base/common/network.ts

@@ -53,6 +53,8 @@ export namespace Schemas {
 	export const vscodeRemoteResource = 'vscode-remote-resource';
 
 	export const userData = 'vscode-userdata';
+
+	export const vscodeCustomEditor = 'vscode-custom-editor';
 }
 
 class RemoteAuthoritiesImpl {

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

@@ -11,6 +11,7 @@ import { Disposable, DisposableStore, dispose, IDisposable, IReference } from 'v
 import { Schemas } from 'vs/base/common/network';
 import { basename } from 'vs/base/common/path';
 import { isWeb } from 'vs/base/common/platform';
+import { isEqual } from 'vs/base/common/resources';
 import { escape } from 'vs/base/common/strings';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import * as modes from 'vs/editor/common/modes';
@@ -28,6 +29,7 @@ import { editorGroupToViewColumn, EditorViewColumn, viewColumnToEditorGroup } fr
 import { IEditorInput, IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
 import { DiffEditorInput } from 'vs/workbench/common/editor/diffEditorInput';
 import { CustomEditorInput } from 'vs/workbench/contrib/customEditor/browser/customEditorInput';
+import { CustomDocumentBackupData } from 'vs/workbench/contrib/customEditor/browser/customEditorInputFactory';
 import { ICustomEditorModel, ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
 import { CustomTextEditorModel } from 'vs/workbench/contrib/customEditor/common/customTextEditorModel';
 import { WebviewExtensionDescription, WebviewIcons } from 'vs/workbench/contrib/webview/browser/webview';
@@ -70,6 +72,10 @@ class WebviewInputStore {
 	public get size(): number {
 		return this._handlesToInputs.size;
 	}
+
+	[Symbol.iterator](): Iterator<WebviewInput> {
+		return this._handlesToInputs.values();
+	}
 }
 
 class WebviewViewTypeTransformer {
@@ -374,7 +380,10 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 
 		const model = modelType === ModelType.Text
 			? CustomTextEditorModel.create(this._instantiationService, viewType, resource)
-			: MainThreadCustomEditorModel.create(this._instantiationService, this._proxy, viewType, resource, cancellation);
+			: MainThreadCustomEditorModel.create(this._instantiationService, this._proxy, viewType, resource, () => {
+				return Array.from(this._webviewInputs)
+					.filter(editor => editor instanceof CustomEditorInput && isEqual(editor.resource, resource)) as CustomEditorInput[];
+			}, cancellation);
 
 		return this._customEditorService.models.add(resource, viewType, model);
 	}
@@ -548,7 +557,6 @@ namespace HotExitState {
 	export type State = typeof Allowed | typeof NotAllowed | Pending;
 }
 
-const customDocumentFileScheme = 'custom';
 
 class MainThreadCustomEditorModel extends Disposable implements ICustomEditorModel, IWorkingCopy {
 
@@ -562,17 +570,19 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		proxy: extHostProtocol.ExtHostWebviewsShape,
 		viewType: string,
 		resource: URI,
+		getEditors: () => CustomEditorInput[],
 		cancellation: CancellationToken,
 	) {
 		const { editable } = await proxy.$createWebviewCustomEditorDocument(resource, viewType, cancellation);
-		return instantiationService.createInstance(MainThreadCustomEditorModel, proxy, viewType, resource, editable);
+		return instantiationService.createInstance(MainThreadCustomEditorModel, proxy, viewType, resource, editable, getEditors);
 	}
 
 	constructor(
 		private readonly _proxy: extHostProtocol.ExtHostWebviewsShape,
 		private readonly _viewType: string,
-		private readonly _realResource: URI,
+		private readonly _editorResource: URI,
 		private readonly _editable: boolean,
+		private readonly _getEditors: () => CustomEditorInput[],
 		@IWorkingCopyService workingCopyService: IWorkingCopyService,
 		@ILabelService private readonly _labelService: ILabelService,
 		@IFileService private readonly _fileService: IFileService,
@@ -587,9 +597,9 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 
 	dispose() {
 		if (this._editable) {
-			this._undoService.removeElements(this._realResource);
+			this._undoService.removeElements(this._editorResource);
 		}
-		this._proxy.$disposeWebviewCustomEditorDocument(this._realResource, this._viewType);
+		this._proxy.$disposeWebviewCustomEditorDocument(this._editorResource, this._viewType);
 		super.dispose();
 	}
 
@@ -598,15 +608,15 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 	public get resource() {
 		// Make sure each custom editor has a unique resource for backup and edits
 		return URI.from({
-			scheme: customDocumentFileScheme,
+			scheme: Schemas.vscodeCustomEditor,
 			authority: this._viewType,
-			path: this._realResource.path,
-			query: JSON.stringify(this._realResource.toJSON())
+			path: this._editorResource.path,
+			query: JSON.stringify(this._editorResource.toJSON()),
 		});
 	}
 
 	public get name() {
-		return basename(this._labelService.getUriLabel(this._realResource));
+		return basename(this._labelService.getUriLabel(this._editorResource));
 	}
 
 	public get capabilities(): WorkingCopyCapabilities {
@@ -645,7 +655,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 
 		this._undoService.pushElement({
 			type: UndoRedoElementType.Resource,
-			resource: this._realResource,
+			resource: this._editorResource,
 			label: label ?? localize('defaultEditLabel', "Edit"),
 			undo: () => this.undo(),
 			redo: () => this.redo(),
@@ -663,7 +673,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		}
 
 		const undoneEdit = this._edits[this._currentEditIndex];
-		await this._proxy.$undo(this._realResource, this.viewType, undoneEdit, this.getEditState());
+		await this._proxy.$undo(this._editorResource, this.viewType, undoneEdit, this.getEditState());
 
 		this.change(() => {
 			--this._currentEditIndex;
@@ -689,7 +699,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		}
 
 		const redoneEdit = this._edits[this._currentEditIndex + 1];
-		await this._proxy.$redo(this._realResource, this.viewType, redoneEdit, this.getEditState());
+		await this._proxy.$redo(this._editorResource, this.viewType, redoneEdit, this.getEditState());
 		this.change(() => {
 			++this._currentEditIndex;
 		});
@@ -704,7 +714,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 			: this._edits.splice(start, toRemove);
 
 		if (removedEdits.length) {
-			this._proxy.$disposeEdits(this._realResource, this._viewType, removedEdits);
+			this._proxy.$disposeEdits(this._editorResource, this._viewType, removedEdits);
 		}
 	}
 
@@ -736,7 +746,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 			editsToRedo = this._edits.slice(this._currentEditIndex, this._savePoint);
 		}
 
-		this._proxy.$revert(this._realResource, this.viewType, { undoneEdits: editsToUndo, redoneEdits: editsToRedo }, this.getEditState());
+		this._proxy.$revert(this._editorResource, this.viewType, { undoneEdits: editsToUndo, redoneEdits: editsToRedo }, this.getEditState());
 		this.change(() => {
 			this._currentEditIndex = this._savePoint;
 			this.spliceEdits();
@@ -747,7 +757,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		if (!this._editable) {
 			return false;
 		}
-		await createCancelablePromise(token => this._proxy.$onSave(this._realResource, this.viewType, token));
+		await createCancelablePromise(token => this._proxy.$onSave(this._editorResource, this.viewType, token));
 		this.change(() => {
 			this._savePoint = this._currentEditIndex;
 		});
@@ -756,7 +766,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 
 	public async saveAs(resource: URI, targetResource: URI, _options?: ISaveOptions): Promise<boolean> {
 		if (this._editable) {
-			await this._proxy.$onSaveAs(this._realResource, this.viewType, targetResource);
+			await this._proxy.$onSaveAs(this._editorResource, this.viewType, targetResource);
 			this.change(() => {
 				this._savePoint = this._currentEditIndex;
 			});
@@ -769,9 +779,25 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 	}
 
 	public async backup(): Promise<IWorkingCopyBackup> {
-		const backupData: IWorkingCopyBackup = {
+		const editors = this._getEditors();
+		if (!editors.length) {
+			throw new Error('No editors found for resource, cannot back up');
+		}
+		const primaryEditor = editors[0];
+
+		const backupData: IWorkingCopyBackup<CustomDocumentBackupData> = {
 			meta: {
 				viewType: this.viewType,
+				editorResource: this._editorResource,
+				extension: primaryEditor.extension ? {
+					id: primaryEditor.extension.id.value,
+					location: primaryEditor.extension.location,
+				} : undefined,
+				webview: {
+					id: primaryEditor.id,
+					options: primaryEditor.webview.options,
+					state: primaryEditor.webview.state,
+				}
 			}
 		};
 
@@ -785,7 +811,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 
 		const pendingState = new HotExitState.Pending(
 			createCancelablePromise(token =>
-				this._proxy.$backup(this._realResource.toJSON(), this.viewType, token)));
+				this._proxy.$backup(this._editorResource.toJSON(), this.viewType, token)));
 		this._hotExitState = pendingState;
 
 		try {

src/vs/workbench/common/editor.ts

@@ -168,6 +168,10 @@ export interface IFileEditorInputFactory {
 	isFileEditorInput(obj: unknown): obj is IFileEditorInput;
 }
 
+interface ICustomEditorInputFactory {
+	createCustomEditorInput(resource: URI, instantiationService: IInstantiationService): Promise<IEditorInput>;
+}
+
 export interface IEditorInputFactoryRegistry {
 
 	/**
@@ -180,6 +184,16 @@ export interface IEditorInputFactoryRegistry {
 	 */
 	getFileEditorInputFactory(): IFileEditorInputFactory;
 
+	/**
+	 * Registers the custom editor input factory to use for custom inputs.
+	 */
+	registerCustomEditorInputFactory(factory: ICustomEditorInputFactory): void;
+
+	/**
+	 * Returns the custom editor input factory to use for custom inputs.
+	 */
+	getCustomEditorInputFactory(): ICustomEditorInputFactory;
+
 	/**
 	 * Registers a editor input factory for the given editor input to the registry. An editor input factory
 	 * is capable of serializing and deserializing editor inputs from string data.
@@ -1387,6 +1401,7 @@ export interface IEditorMemento<T> {
 class EditorInputFactoryRegistry implements IEditorInputFactoryRegistry {
 	private instantiationService: IInstantiationService | undefined;
 	private fileEditorInputFactory: IFileEditorInputFactory | undefined;
+	private customEditorInputFactory: ICustomEditorInputFactory | undefined;
 
 	private readonly editorInputFactoryConstructors: Map<string, IConstructorSignature0<IEditorInputFactory>> = new Map();
 	private readonly editorInputFactoryInstances: Map<string, IEditorInputFactory> = new Map();
@@ -1414,6 +1429,14 @@ class EditorInputFactoryRegistry implements IEditorInputFactoryRegistry {
 		return assertIsDefined(this.fileEditorInputFactory);
 	}
 
+	registerCustomEditorInputFactory(factory: ICustomEditorInputFactory): void {
+		this.customEditorInputFactory = factory;
+	}
+
+	getCustomEditorInputFactory(): ICustomEditorInputFactory {
+		return assertIsDefined(this.customEditorInputFactory);
+	}
+
 	registerEditorInputFactory(editorInputId: string, ctor: IConstructorSignature0<IEditorInputFactory>): IDisposable {
 		if (!this.instantiationService) {
 			this.editorInputFactoryConstructors.set(editorInputId, ctor);

src/vs/workbench/contrib/backup/common/backupRestorer.ts

@@ -10,9 +10,11 @@ import { IEditorService } from 'vs/workbench/services/editor/common/editorServic
 import { IResourceEditorInput } from 'vs/platform/editor/common/editor';
 import { Schemas } from 'vs/base/common/network';
 import { ILifecycleService, LifecyclePhase } from 'vs/platform/lifecycle/common/lifecycle';
-import { IUntitledTextResourceEditorInput, IEditorInput } from 'vs/workbench/common/editor';
+import { IUntitledTextResourceEditorInput, IEditorInput, IEditorInputFactoryRegistry, Extensions as EditorExtensions, IEditorInputWithOptions } from 'vs/workbench/common/editor';
 import { toLocalResource, isEqual } from 'vs/base/common/resources';
 import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/common/environmentService';
+import { Registry } from 'vs/platform/registry/common/platform';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 
 export class BackupRestorer implements IWorkbenchContribution {
 
@@ -22,7 +24,8 @@ export class BackupRestorer implements IWorkbenchContribution {
 		@IEditorService private readonly editorService: IEditorService,
 		@IBackupFileService private readonly backupFileService: IBackupFileService,
 		@ILifecycleService private readonly lifecycleService: ILifecycleService,
-		@IWorkbenchEnvironmentService private readonly environmentService: IWorkbenchEnvironmentService
+		@IWorkbenchEnvironmentService private readonly environmentService: IWorkbenchEnvironmentService,
+		@IInstantiationService private readonly instantiationService: IInstantiationService,
 	) {
 		this.restoreBackups();
 	}
@@ -78,13 +81,13 @@ export class BackupRestorer implements IWorkbenchContribution {
 
 	private async doOpenEditors(resources: URI[]): Promise<void> {
 		const hasOpenedEditors = this.editorService.visibleEditors.length > 0;
-		const inputs = resources.map((resource, index) => this.resolveInput(resource, index, hasOpenedEditors));
+		const inputs = await Promise.all(resources.map((resource, index) => this.resolveInput(resource, index, hasOpenedEditors)));
 
 		// Open all remaining backups as editors and resolve them to load their backups
 		await this.editorService.openEditors(inputs);
 	}
 
-	private resolveInput(resource: URI, index: number, hasOpenedEditors: boolean): IResourceEditorInput | IUntitledTextResourceEditorInput {
+	private async resolveInput(resource: URI, index: number, hasOpenedEditors: boolean): Promise<IResourceEditorInput | IUntitledTextResourceEditorInput | IEditorInputWithOptions> {
 		const options = { pinned: true, preserveFocus: true, inactive: index > 0 || hasOpenedEditors };
 
 		// this is a (weak) strategy to find out if the untitled input had
@@ -94,6 +97,12 @@ export class BackupRestorer implements IWorkbenchContribution {
 			return { resource: toLocalResource(resource, this.environmentService.configuration.remoteAuthority), options, forceUntitled: true };
 		}
 
+		if (resource.scheme === Schemas.vscodeCustomEditor) {
+			const editor = await Registry.as<IEditorInputFactoryRegistry>(EditorExtensions.EditorInputFactories).getCustomEditorInputFactory()
+				.createCustomEditorInput(resource, this.instantiationService);
+			return { editor, options };
+		}
+
 		return { resource, options };
 	}
 }

src/vs/workbench/contrib/customEditor/browser/customEditorInputFactory.ts

@@ -3,21 +3,41 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { URI } from 'vs/base/common/uri';
+import { Lazy } from 'vs/base/common/lazy';
+import { URI, UriComponents } from 'vs/base/common/uri';
 import { generateUuid } from 'vs/base/common/uuid';
+import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { IEditorInput } from 'vs/workbench/common/editor';
 import { CustomEditorInput } from 'vs/workbench/contrib/customEditor/browser/customEditorInput';
+import { IWebviewService } from 'vs/workbench/contrib/webview/browser/webview';
 import { WebviewEditorInputFactory } from 'vs/workbench/contrib/webview/browser/webviewEditorInputFactory';
-import { IWebviewWorkbenchService } from 'vs/workbench/contrib/webview/browser/webviewWorkbenchService';
-import { Lazy } from 'vs/base/common/lazy';
+import { IWebviewWorkbenchService, WebviewInputOptions } from 'vs/workbench/contrib/webview/browser/webviewWorkbenchService';
+import { IBackupFileService } from 'vs/workbench/services/backup/common/backup';
+
+export interface CustomDocumentBackupData {
+	readonly viewType: string;
+	readonly editorResource: UriComponents;
+	readonly extension: undefined | {
+		readonly location: UriComponents;
+		readonly id: string;
+	};
+
+	readonly webview: {
+		readonly id: string;
+		readonly options: WebviewInputOptions;
+		readonly state: any;
+	};
+}
 
 export class CustomEditorInputFactory extends WebviewEditorInputFactory {
 
 	public static readonly ID = CustomEditorInput.typeId;
 
 	public constructor(
+		@IWebviewWorkbenchService webviewWorkbenchService: IWebviewWorkbenchService,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
-		@IWebviewWorkbenchService private readonly webviewWorkbenchService: IWebviewWorkbenchService,
+		@IWebviewService private readonly _webviewService: IWebviewService,
 	) {
 		super(webviewWorkbenchService);
 	}
@@ -43,11 +63,19 @@ export class CustomEditorInputFactory extends WebviewEditorInputFactory {
 		const id = data.id || generateUuid();
 
 		const webview = new Lazy(() => {
-			const webviewInput = this.webviewWorkbenchService.reviveWebview(id, data.viewType, data.title, data.iconPath, data.state, data.options, data.extensionLocation && data.extensionId ? {
-				location: data.extensionLocation,
-				id: data.extensionId
-			} : undefined, data.group);
-			return webviewInput.webview;
+			const webview = this._webviewService.createWebviewOverlay(id, {
+				enableFindWidget: data.options.enableFindWidget,
+				retainContextWhenHidden: data.options.retainContextWhenHidden
+			}, data.options);
+
+			if (data.extensionLocation && data.extensionId) {
+				webview.extension = {
+					location: data.extensionLocation,
+					id: data.extensionId
+				};
+			}
+
+			return webview;
 		});
 
 		const customInput = this._instantiationService.createInstance(CustomEditorInput, URI.from((data as any).editorResource), data.viewType, id, webview);
@@ -56,4 +84,37 @@ export class CustomEditorInputFactory extends WebviewEditorInputFactory {
 		}
 		return customInput;
 	}
+
+	public static createCustomEditorInput(resource: URI, instantiationService: IInstantiationService): Promise<IEditorInput> {
+		return instantiationService.invokeFunction(async accessor => {
+			const webviewService = accessor.get<IWebviewService>(IWebviewService);
+			const backupFileService = accessor.get<IBackupFileService>(IBackupFileService);
+
+			const backup = await backupFileService.resolve(resource);
+			if (!backup) {
+				throw new Error(`No backup found for custom editor: ${resource}`);
+			}
+
+			const backupData = backup.meta as CustomDocumentBackupData;
+			const id = backupData.webview.id;
+
+			const webview = new Lazy(() => {
+				const webview = webviewService.createWebviewOverlay(id, {
+					enableFindWidget: backupData.webview.options.enableFindWidget,
+					retainContextWhenHidden: backupData.webview.options.retainContextWhenHidden
+				}, backupData.webview.options);
+
+				webview.extension = backupData.extension ? {
+					location: URI.revive(backupData.extension.location),
+					id: new ExtensionIdentifier(backupData.extension.id),
+				} : undefined;
+
+				return webview;
+			});
+
+			const editor = instantiationService.createInstance(CustomEditorInput, URI.revive(backupData.editorResource), backupData.viewType, id, webview);
+			editor.updateGroup(0);
+			return editor;
+		});
+	}
 }

src/vs/workbench/contrib/customEditor/browser/webviewEditor.contribution.ts

@@ -38,6 +38,8 @@ Registry.as<IEditorInputFactoryRegistry>(EditorInputExtensions.EditorInputFactor
 	CustomEditorInputFactory.ID,
 	CustomEditorInputFactory);
 
+Registry.as<IEditorInputFactoryRegistry>(EditorInputExtensions.EditorInputFactories).registerCustomEditorInputFactory(CustomEditorInputFactory);
+
 Registry.as<IConfigurationRegistry>(ConfigurationExtensions.Configuration)
 	.registerConfiguration({
 		...workbenchConfigurationNodeBase,

src/vs/workbench/services/workingCopy/common/workingCopyService.ts

@@ -27,12 +27,12 @@ export const enum WorkingCopyCapabilities {
  * `IBackupFileService.resolve(workingCopy.resource)` to
  * retrieve the backup when loading the working copy.
  */
-export interface IWorkingCopyBackup {
+export interface IWorkingCopyBackup<MetaType = object> {
 
 	/**
 	 * Any serializable metadata to be associated with the backup.
 	 */
-	meta?: object;
+	meta?: MetaType;
 
 	/**
 	 * Use this for larger textual content of the backup.