Use the UndoRedoService for CustomEditors (#92408)

* Use the UndoRedoService for CustomEditors

For #90110

Changes custom editors (the ones not based on text) to use the UndoRedoService. This involved:

- Moving edit lifecycle back into the main process again (this is actually the bulk of the changes)
- Removing the `undo`/`redo` methods on  `CustomEditorModel`
- Using the undoRedoService to trigger undo/redo

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

@@ -6,7 +6,7 @@
 import { CancelablePromise, createCancelablePromise } from 'vs/base/common/async';
 import { onUnexpectedError } from 'vs/base/common/errors';
 import { Emitter, Event } from 'vs/base/common/event';
-import { Disposable, DisposableStore, IDisposable, IReference, dispose } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, dispose, IDisposable, IReference } from 'vs/base/common/lifecycle';
 import { Schemas } from 'vs/base/common/network';
 import { basename } from 'vs/base/common/path';
 import { isWeb } from 'vs/base/common/platform';
@@ -21,6 +21,7 @@ import { ILabelService } from 'vs/platform/label/common/label';
 import { IOpenerService } from 'vs/platform/opener/common/opener';
 import { IProductService } from 'vs/platform/product/common/productService';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
+import { IUndoRedoService, UndoRedoElementType } from 'vs/platform/undoRedo/common/undoRedo';
 import * as extHostProtocol from 'vs/workbench/api/common/extHost.protocol';
 import { editorGroupToViewColumn, EditorViewColumn, viewColumnToEditorGroup } from 'vs/workbench/api/common/shared/editor';
 import { IEditorInput, IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
@@ -365,12 +366,14 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 		return this._customEditorService.models.add(resource, viewType, model);
 	}
 
-	public async $onDidChangeCustomDocumentState(resource: UriComponents, viewType: string, state: { dirty: boolean }) {
-		const model = await this._customEditorService.models.get(URI.revive(resource), viewType);
+	public async $onDidEdit(resourceComponents: UriComponents, viewType: string, editId: number): Promise<void> {
+		const resource = URI.revive(resourceComponents);
+		const model = await this._customEditorService.models.get(resource, viewType);
 		if (!model || !(model instanceof MainThreadCustomEditorModel)) {
 			throw new Error('Could not find model for webview editor');
 		}
-		model.setDirty(state.dirty);
+
+		model.pushEdit(editId);
 	}
 
 	private hookupWebviewEventDelegate(handle: extHostProtocol.WebviewPanelHandle, input: WebviewInput) {
@@ -536,7 +539,9 @@ namespace HotExitState {
 class MainThreadCustomEditorModel extends Disposable implements ICustomEditorModel, IWorkingCopy {
 
 	private _hotExitState: HotExitState.State = HotExitState.Allowed;
-	private _dirty = false;
+	private _currentEditIndex: number = -1;
+	private _savePoint: number = -1;
+	private readonly _edits: Array<number> = [];
 
 	public static async create(
 		instantiationService: IInstantiationService,
@@ -556,19 +561,25 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		@IWorkingCopyService workingCopyService: IWorkingCopyService,
 		@ILabelService private readonly _labelService: ILabelService,
 		@IFileService private readonly _fileService: IFileService,
+		@IUndoRedoService private readonly _undoService: IUndoRedoService,
 	) {
 		super();
-		this._register(workingCopyService.registerWorkingCopy(this));
+		if (_editable) {
+			this._register(workingCopyService.registerWorkingCopy(this));
+		}
 	}
 
 	dispose() {
+		if (this._editable) {
+			this._undoService.removeElements(this.resource);
+		}
 		this._proxy.$disposeWebviewCustomEditorDocument(this.resource, this._viewType);
 		super.dispose();
 	}
 
 	//#region IWorkingCopy
 
-	public get resource() { return this._resource; }
+	public get resource() { return this._resource; } // custom://viewType/path/file
 
 	public get name() {
 		return basename(this._labelService.getUriLabel(this._resource));
@@ -579,7 +590,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 	}
 
 	public isDirty(): boolean {
-		return this._dirty;
+		return this._edits.length > 0 && this._savePoint !== this._currentEditIndex;
 	}
 
 	private readonly _onDidChangeDirty: Emitter<void> = this._register(new Emitter<void>());
@@ -589,36 +600,106 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 	readonly onDidChangeContent: Event<void> = this._onDidChangeContent.event;
 
 	//#endregion
-
 	public get viewType() {
 		return this._viewType;
 	}
 
-	public setDirty(dirty: boolean): void {
+	public pushEdit(editId: number) {
+		if (!this._editable) {
+			throw new Error('Document is not editable');
+		}
+
+		this.change(() => {
+			this.spliceEdits(editId);
+			this._currentEditIndex = this._edits.length - 1;
+		});
+
+		this._undoService.pushElement({
+			type: UndoRedoElementType.Resource,
+			resource: this.resource,
+			label: 'Edit', // TODO: get this from extensions?
+			undo: async () => {
+				if (!this._editable) {
+					return;
+				}
+
+				if (this._currentEditIndex < 0) {
+					// nothing to undo
+					return;
+				}
+
+				const undoneEdit = this._edits[this._currentEditIndex];
+				await this._proxy.$undo(this.resource, this.viewType, undoneEdit);
+
+				this.change(() => {
+					--this._currentEditIndex;
+				});
+			},
+			redo: async () => {
+				if (!this._editable) {
+					return;
+				}
+
+				if (this._currentEditIndex >= this._edits.length - 1) {
+					// nothing to redo
+					return;
+				}
+
+				const redoneEdit = this._edits[this._currentEditIndex + 1];
+				await this._proxy.$redo(this.resource, this.viewType, redoneEdit);
+				this.change(() => {
+					++this._currentEditIndex;
+				});
+			}
+		});
+	}
+
+	private spliceEdits(editToInsert?: number) {
+		const start = this._currentEditIndex + 1;
+		const toRemove = this._edits.length - this._currentEditIndex;
+
+		const removedEdits = typeof editToInsert === 'number'
+			? this._edits.splice(start, toRemove, editToInsert)
+			: this._edits.splice(start, toRemove);
+
+		if (removedEdits.length) {
+			this._proxy.$disposeEdits(this.resource, this._viewType, removedEdits);
+		}
+	}
+
+	private change(makeEdit: () => void): void {
+		const wasDirty = this.isDirty();
+		makeEdit();
 		this._onDidChangeContent.fire();
 
-		if (this._dirty !== dirty) {
-			this._dirty = dirty;
+		if (this.isDirty() !== wasDirty) {
 			this._onDidChangeDirty.fire();
 		}
 	}
 
 	public async revert(_options?: IRevertOptions) {
-		if (this._editable) {
-			this._proxy.$revert(this.resource, this.viewType);
+		if (!this._editable) {
+			return;
 		}
-	}
 
-	public undo() {
-		if (this._editable) {
-			this._proxy.$undo(this.resource, this.viewType);
+		if (this._currentEditIndex === this._savePoint) {
+			return;
 		}
-	}
 
-	public redo() {
-		if (this._editable) {
-			this._proxy.$redo(this.resource, this.viewType);
+		let editsToUndo: number[] = [];
+		let editsToRedo: number[] = [];
+
+		if (this._currentEditIndex >= this._savePoint) {
+			editsToUndo = this._edits.slice(this._savePoint, this._currentEditIndex).reverse();
+		} else if (this._currentEditIndex < this._savePoint) {
+			editsToRedo = this._edits.slice(this._currentEditIndex, this._savePoint);
 		}
+
+		this._proxy.$revert(this.resource, this.viewType, { undoneEdits: editsToUndo, redoneEdits: editsToRedo });
+		this.change(() => {
+			this._currentEditIndex = this._savePoint;
+			this.spliceEdits();
+		});
 	}
 
 	public async save(_options?: ISaveOptions): Promise<boolean> {
@@ -626,14 +707,18 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 			return false;
 		}
 		await createCancelablePromise(token => this._proxy.$onSave(this.resource, this.viewType, token));
-		this.setDirty(false);
+		this.change(() => {
+			this._savePoint = this._currentEditIndex;
+		});
 		return true;
 	}
 
 	public async saveAs(resource: URI, targetResource: URI, _options?: ISaveOptions): Promise<boolean> {
 		if (this._editable) {
 			await this._proxy.$onSaveAs(this.resource, this.viewType, targetResource);
-			this.setDirty(false);
+			this.change(() => {
+				this._savePoint = this._currentEditIndex;
+			});
 			return true;
 		} else {
 			// Since the editor is readonly, just copy the file over

src/vs/workbench/api/common/extHost.protocol.ts

@@ -592,7 +592,7 @@ export interface MainThreadWebviewsShape extends IDisposable {
 	$registerCustomEditorProvider(extension: WebviewExtensionDescription, viewType: string, options: modes.IWebviewPanelOptions): void;
 	$unregisterEditorProvider(viewType: string): void;
 
-	$onDidChangeCustomDocumentState(resource: UriComponents, viewType: string, state: { dirty: boolean }): void;
+	$onDidEdit(resource: UriComponents, viewType: string, editId: number): void;
 }
 
 export interface WebviewPanelViewStateData {
@@ -615,9 +615,11 @@ export interface ExtHostWebviewsShape {
 	$createWebviewCustomEditorDocument(resource: UriComponents, viewType: string): Promise<{ editable: boolean }>;
 	$disposeWebviewCustomEditorDocument(resource: UriComponents, viewType: string): Promise<void>;
 
-	$undo(resource: UriComponents, viewType: string): void;
-	$redo(resource: UriComponents, viewType: string): void;
-	$revert(resource: UriComponents, viewType: string): void;
+	$undo(resource: UriComponents, viewType: string, editId: number): Promise<void>;
+	$redo(resource: UriComponents, viewType: string, editId: number): Promise<void>;
+	$revert(resource: UriComponents, viewType: string, changes: { undoneEdits: number[], redoneEdits: number[] }): Promise<void>;
+	$disposeEdits(resourceComponents: UriComponents, viewType: string, editIds: number[]): void;
+
 	$onSave(resource: UriComponents, viewType: string, cancellation: CancellationToken): Promise<void>;
 	$onSaveAs(resource: UriComponents, viewType: string, targetResource: UriComponents): Promise<void>;
 

src/vs/workbench/api/common/extHostWebview.ts

@@ -18,6 +18,7 @@ import { IExtHostWorkspace } from 'vs/workbench/api/common/extHostWorkspace';
 import { EditorViewColumn } from 'vs/workbench/api/common/shared/editor';
 import { asWebviewUri, WebviewInitData } from 'vs/workbench/api/common/shared/webview';
 import type * as vscode from 'vscode';
+import { Cache } from './cache';
 import { ExtHostWebviewsShape, IMainContext, MainContext, MainThreadWebviewsShape, WebviewExtensionDescription, WebviewPanelHandle, WebviewPanelViewStateData } from './extHost.protocol';
 import { Disposable as VSCodeDisposable } from './extHostTypes';
 
@@ -245,8 +246,6 @@ export class ExtHostWebviewEditor extends Disposable implements vscode.WebviewPa
 	}
 }
 
-type EditType = unknown;
-
 class CustomDocument extends Disposable implements vscode.CustomDocument {
 
 	public static create(proxy: MainThreadWebviewsShape, viewType: string, uri: vscode.Uri) {
@@ -255,9 +254,7 @@ class CustomDocument extends Disposable implements vscode.CustomDocument {
 
 	// Explicitly initialize all properties as we seal the object after creation!
 
-	#currentEditIndex: number = -1;
-	#savePoint: number = -1;
-	readonly #edits: Array<EditType> = [];
+	readonly #_edits = new Cache<unknown>('edits');
 
 	readonly #proxy: MainThreadWebviewsShape;
 	readonly #viewType: string;
@@ -299,58 +296,28 @@ class CustomDocument extends Disposable implements vscode.CustomDocument {
 
 		this.#capabilities = capabilities;
 		capabilities.editing?.onDidEdit(edit => {
-			this.pushEdit(edit);
+			const id = this.#_edits.add([edit]);
+			this.#proxy.$onDidEdit(this.uri, this.viewType, id);
 		});
 	}
 
-	/** @internal*/ async _revert() {
+	/** @internal*/ async _revert(changes: { undoneEdits: number[], redoneEdits: number[] }) {
 		const editing = this.getEditingCapability();
-		if (this.#currentEditIndex === this.#savePoint) {
-			return true;
-		}
-
-
-		let undoneEdits: EditType[] = [];
-		let appliedEdits: EditType[] = [];
-		if (this.#currentEditIndex >= this.#savePoint) {
-			undoneEdits = this.#edits.slice(this.#savePoint, this.#currentEditIndex).reverse();
-		} else if (this.#currentEditIndex < this.#savePoint) {
-			appliedEdits = this.#edits.slice(this.#currentEditIndex, this.#savePoint);
-		}
-
-		this.#currentEditIndex = this.#savePoint;
-		this.spliceEdits();
-
-		await editing.revert({ undoneEdits, appliedEdits });
-
-		this.updateState();
-		return true;
+		const undoneEdits = changes.undoneEdits.map(id => this.#_edits.get(id, 0));
+		const appliedEdits = changes.redoneEdits.map(id => this.#_edits.get(id, 0));
+		return editing.revert({ undoneEdits, appliedEdits });
 	}
 
-	/** @internal*/ _undo() {
+	/** @internal*/ _undo(editId: number) {
 		const editing = this.getEditingCapability();
-		if (this.#currentEditIndex < 0) {
-			// nothing to undo
-			return;
-		}
-
-		const undoneEdit = this.#edits[this.#currentEditIndex];
-		--this.#currentEditIndex;
-		editing.undoEdits([undoneEdit]);
-		this.updateState();
+		const edit = this.#_edits.get(editId, 0);
+		return editing.undoEdits([edit]);
 	}
 
-	/** @internal*/ _redo() {
+	/** @internal*/ _redo(editId: number) {
 		const editing = this.getEditingCapability();
-		if (this.#currentEditIndex >= this.#edits.length - 1) {
-			// nothing to redo
-			return;
-		}
-
-		++this.#currentEditIndex;
-		const redoneEdit = this.#edits[this.#currentEditIndex];
-		editing.applyEdits([redoneEdit]);
-		this.updateState();
+		const edit = this.#_edits.get(editId, 0);
+		return editing.applyEdits([edit]);
 	}
 
 	/** @internal*/ _save(cancellation: CancellationToken) {
@@ -365,29 +332,14 @@ class CustomDocument extends Disposable implements vscode.CustomDocument {
 		return this.getEditingCapability().backup(cancellation);
 	}
 
+	/** @internal*/ _disposeEdits(editIds: number[]) {
+		for (const editId of editIds) {
+			this.#_edits.delete(editId);
+		}
+	}
+
 	//#endregion
 
-	private pushEdit(edit: EditType) {
-		this.spliceEdits(edit);
-
-		this.#currentEditIndex = this.#edits.length - 1;
-		this.updateState();
-	}
-
-	private updateState() {
-		const dirty = this.#edits.length > 0 && this.#savePoint !== this.#currentEditIndex;
-		this.#proxy.$onDidChangeCustomDocumentState(this.uri, this.viewType, { dirty });
-	}
-
-	private spliceEdits(editToInsert?: EditType) {
-		const start = this.#currentEditIndex + 1;
-		const toRemove = this.#edits.length - this.#currentEditIndex;
-
-		editToInsert
-			? this.#edits.splice(start, toRemove, editToInsert)
-			: this.#edits.splice(start, toRemove);
-	}
-
 	private getEditingCapability(): vscode.CustomEditorEditingCapability {
 		if (!this.#capabilities?.editing) {
 			throw new Error('Document is not editable');
@@ -702,24 +654,29 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		}
 	}
 
-	async $undo(resourceComponents: UriComponents, viewType: string): Promise<void> {
+	$disposeEdits(resourceComponents: UriComponents, viewType: string, editIds: number[]): void {
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		document._undo();
+		document._disposeEdits(editIds);
 	}
 
-	async $redo(resourceComponents: UriComponents, viewType: string): Promise<void> {
+	async $undo(resourceComponents: UriComponents, viewType: string, editId: number): Promise<void> {
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		document._redo();
+		return document._undo(editId);
 	}
 
-	async $revert(resourceComponents: UriComponents, viewType: string): Promise<void> {
+	async $redo(resourceComponents: UriComponents, viewType: string, editId: number): Promise<void> {
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		document._revert();
+		return document._redo(editId);
+	}
+
+	async $revert(resourceComponents: UriComponents, viewType: string, changes: { undoneEdits: number[], redoneEdits: number[] }): Promise<void> {
+		const document = this.getCustomDocument(viewType, resourceComponents);
+		return document._revert(changes);
 	}
 
 	async $onSave(resourceComponents: UriComponents, viewType: string, cancellation: CancellationToken): Promise<void> {
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		document._save(cancellation);
+		return document._save(cancellation);
 	}
 
 	async $onSaveAs(resourceComponents: UriComponents, viewType: string, targetResource: UriComponents): Promise<void> {

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

@@ -15,6 +15,7 @@ import { IFileDialogService } from 'vs/platform/dialogs/common/dialogs';
 import { ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ILabelService } from 'vs/platform/label/common/label';
+import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
 import { GroupIdentifier, IEditorInput, IRevertOptions, ISaveOptions, Verbosity } from 'vs/workbench/common/editor';
 import { ICustomEditorModel, ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
 import { IWebviewService, WebviewOverlay } from 'vs/workbench/contrib/webview/browser/webview';
@@ -44,6 +45,7 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 		@IFileDialogService private readonly fileDialogService: IFileDialogService,
 		@IFilesConfigurationService private readonly filesConfigurationService: IFilesConfigurationService,
 		@IEditorService private readonly editorService: IEditorService,
+		@IUndoRedoService private readonly undoRedoService: IUndoRedoService,
 	) {
 		super(id, viewType, '', webview, webviewService, webviewWorkbenchService);
 		this._editorResource = resource;
@@ -175,10 +177,12 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 	}
 
 	public undo(): void {
-		assertIsDefined(this._modelRef).object.undo();
+		assertIsDefined(this._modelRef);
+		this.undoRedoService.undo(this.resource);
 	}
 
 	public redo(): void {
-		assertIsDefined(this._modelRef).object.redo();
+		assertIsDefined(this._modelRef);
+		this.undoRedoService.redo(this.resource);
 	}
 }

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

@@ -52,8 +52,6 @@ export interface ICustomEditorModel extends IDisposable {
 	isDirty(): boolean;
 	readonly onDidChangeDirty: Event<void>;
 
-	undo(): void;
-	redo(): void;
 	revert(options?: IRevertOptions): Promise<void>;
 
 	save(options?: ISaveOptions): Promise<boolean>;

src/vs/workbench/contrib/customEditor/common/customTextEditorModel.ts

@@ -64,14 +64,6 @@ export class CustomTextEditorModel extends Disposable implements ICustomEditorMo
 		return this.textFileService.revert(this.resource, options);
 	}
 
-	public undo() {
-		this.textFileService.files.get(this.resource)?.textEditorModel?.undo();
-	}
-
-	public redo() {
-		this.textFileService.files.get(this.resource)?.textEditorModel?.redo();
-	}
-
 	public async save(options?: ISaveOptions): Promise<boolean> {
 		return !!await this.textFileService.save(this.resource, options);
 	}