Update custom editor api

For #77131

- Use a class for `CustomDocument` instead of an interface. Extensions can now add their own data to a `CustomDocument` by sublassing

- Renamed `resolveCustomDocument` to `openCustomDocument` and require that extensions return a `CustomDocument`

- Exposed edits on `CustomDocument`

- Made the third parameter of `registerCustomEditorProvider` a generic options bag that takes a `webviewOptions`

extensions/image-preview/src/preview.ts

@@ -27,8 +27,8 @@ export class PreviewManager implements vscode.CustomEditorProvider {
 		private readonly zoomStatusBarEntry: ZoomStatusBarEntry,
 	) { }
 
-	public async resolveCustomDocument(_document: vscode.CustomDocument): Promise<void> {
+	public async openCustomDocument(uri: vscode.Uri) {
-		// noop
+		return new vscode.CustomDocument(PreviewManager.viewType, uri);
 	}
 
 	public async resolveCustomEditor(

extensions/markdown-language-features/src/features/previewManager.ts

@@ -63,6 +63,8 @@ export class MarkdownPreviewManager extends Disposable implements vscode.Webview
 
 	private _activePreview: DynamicMarkdownPreview | undefined = undefined;
 
+	private readonly customEditorViewType = 'vscode.markdown.preview.editor';
+
 	public constructor(
 		private readonly _contentProvider: MarkdownContentProvider,
 		private readonly _logger: Logger,
@@ -70,7 +72,7 @@ export class MarkdownPreviewManager extends Disposable implements vscode.Webview
 	) {
 		super();
 		this._register(vscode.window.registerWebviewPanelSerializer(DynamicMarkdownPreview.viewType, this));
-		this._register(vscode.window.registerCustomEditorProvider('vscode.markdown.preview.editor', this));
+		this._register(vscode.window.registerCustomEditorProvider(this.customEditorViewType, this));
 	}
 
 	public refresh() {
@@ -148,8 +150,8 @@ export class MarkdownPreviewManager extends Disposable implements vscode.Webview
 		this.registerDynamicPreview(preview);
 	}
 
-	public async resolveCustomDocument(_document: vscode.CustomDocument): Promise<void> {
+	public async openCustomDocument(uri: vscode.Uri) {
-		// noop
+		return new vscode.CustomDocument(this.customEditorViewType, uri);
 	}
 
 	public async resolveCustomTextEditor(

src/vs/vscode.proposed.d.ts

@@ -1289,17 +1289,11 @@ declare module 'vscode' {
 
 	//#region Custom editors: https://github.com/microsoft/vscode/issues/77131
 
-	// TODO:
-	// - Think about where a rename would live.
-	// - Think about handling go to line? (add other editor options? reveal?)
-	// - Should we expose edits?
-	// - More properties from `TextDocument`?
-
 	/**
 	 * Defines the editing capability of a custom webview editor. This allows the webview editor to hook into standard
 	 * editor events such as `undo` or `save`.
 	 *
-	 * @param EditType Type of edits.
+	 * @param EditType Type of edits used for the documents this delegate handles.
 	 */
 	interface CustomEditorEditingDelegate<EditType = unknown> {
 		/**
@@ -1310,7 +1304,7 @@ declare module 'vscode' {
 		 *
 		 * @return Thenable signaling that the save has completed.
 		 */
-		save(document: CustomDocument, cancellation: CancellationToken): Thenable<void>;
+		save(document: CustomDocument<EditType>, cancellation: CancellationToken): Thenable<void>;
 
 		/**
 		 * Save the existing resource at a new path.
@@ -1320,7 +1314,7 @@ declare module 'vscode' {
 		 *
 		 * @return Thenable signaling that the save has completed.
 		 */
-		saveAs(document: CustomDocument, targetResource: Uri): Thenable<void>;
+		saveAs(document: CustomDocument<EditType>, targetResource: Uri): Thenable<void>;
 
 		/**
 		 * Event triggered by extensions to signal to VS Code that an edit has occurred.
@@ -1337,7 +1331,7 @@ declare module 'vscode' {
 		 *
 		 * @return Thenable signaling that the change has completed.
 		 */
-		applyEdits(document: CustomDocument, edits: readonly EditType[]): Thenable<void>;
+		applyEdits(document: CustomDocument<EditType>, edits: readonly EditType[]): Thenable<void>;
 
 		/**
 		 * Undo a set of edits.
@@ -1349,7 +1343,7 @@ declare module 'vscode' {
 		 *
 		 * @return Thenable signaling that the change has completed.
 		 */
-		undoEdits(document: CustomDocument, edits: readonly EditType[]): Thenable<void>;
+		undoEdits(document: CustomDocument<EditType>, edits: readonly EditType[]): Thenable<void>;
 
 		/**
 		 * Revert the file to its last saved state.
@@ -1359,7 +1353,7 @@ declare module 'vscode' {
 		 *
 		 * @return Thenable signaling that the change has completed.
 		 */
-		revert(document: CustomDocument, edits: CustomDocumentRevert<EditType>): Thenable<void>;
+		revert(document: CustomDocument<EditType>, edits: CustomDocumentRevert<EditType>): Thenable<void>;
 
 		/**
 		 * Back up the resource in its current state.
@@ -1380,22 +1374,25 @@ declare module 'vscode' {
 		 * in an operation that takes time to complete, your extension may decide to finish the ongoing backup rather
 		 * than cancelling it to ensure that VS Code has some valid backup.
 		 */
-		backup(document: CustomDocument, cancellation: CancellationToken): Thenable<void>;
+		backup(document: CustomDocument<EditType>, cancellation: CancellationToken): Thenable<void>;
 	}
 
 	/**
-	 * Event triggered by extensions to signal to VS Code that an edit has occurred on a CustomDocument``.
+	 * Event triggered by extensions to signal to VS Code that an edit has occurred on a `CustomDocument`.
+	 *
+	 * @param EditType Type of edits used for the document.
 	 */
 	interface CustomDocumentEditEvent<EditType = unknown> {
 		/**
 		 * Document the edit is for.
 		 */
-		readonly document: CustomDocument;
+		readonly document: CustomDocument<EditType>;
 
 		/**
 		 * Object that describes the edit.
 		 *
-		 * Edit objects are passed back to your extension in `undoEdits`, `applyEdits`, and `revert`.
+		 * Edit objects are passed back to your extension in `CustomEditorEditingDelegate.undoEdits`,
+		 * `CustomEditorEditingDelegate.applyEdits`, and `CustomEditorEditingDelegate.revert`.
 		 */
 		readonly edit: EditType;
 
@@ -1423,13 +1420,19 @@ declare module 'vscode' {
 	/**
 	 * Represents a custom document used by a `CustomEditorProvider`.
 	 *
-	 * Custom documents are only used within a given `CustomEditorProvider`. The lifecycle of a
+	 * All custom documents must subclass `CustomDocument`. Custom documents are only used within a given
-	 * `CustomDocument` is managed by VS Code. When no more references remain to a given `CustomDocument`,
+	 * `CustomEditorProvider`. The lifecycle of a `CustomDocument` is managed by VS Code. When no more references
-	 * then it is disposed of.
+	 * remain to a `CustomDocument`, it is disposed of.
 	 *
-	 * @param UserDataType Type of custom object that extensions can store on the document.
+	 * @param EditType Type of edits used in this document.
 	 */
-	interface CustomDocument<UserDataType = unknown> {
+	class CustomDocument<EditType = unknown> {
+		/**
+		 * @param viewType The associated uri for this document.
+		 * @param uri The associated viewType for this document.
+		 */
+		constructor(viewType: string, uri: Uri);
+
 		/**
 		 * The associated viewType for this document.
 		 */
@@ -1446,12 +1449,17 @@ declare module 'vscode' {
 		readonly onDidDispose: Event<void>;
 
 		/**
-		 * Custom data that an extension can store on the document.
+		 * List of edits from document open to the document's current state.
 		 */
-		userData?: UserDataType;
+		readonly appliedEdits: ReadonlyArray<EditType>;
 
-		// TODO: Should we expose edits here?
+		/**
-		// This could be helpful for tracking the life cycle of edits
+		 * List of edits from document open to the document's last saved point.
+		 *
+		 * The save point will be behind `appliedEdits` if the user saves and then continues editing,
+		 * or in front of the last entry in `appliedEdits` if the user saves and then hits undo.
+		 */
+		readonly savedEdits: ReadonlyArray<EditType>;
 	}
 
 	/**
@@ -1463,7 +1471,8 @@ declare module 'vscode' {
 	 * You should use custom text based editors when dealing with binary files or more complex scenarios. For simple text
 	 * based documents, use [`WebviewTextEditorProvider`](#WebviewTextEditorProvider) instead.
 	 */
-	export interface CustomEditorProvider {
+	export interface CustomEditorProvider<EditType = unknown> {
+
 		/**
 		 * Resolve the model for a given resource.
 		 *
@@ -1472,18 +1481,18 @@ declare module 'vscode' {
 		 * If all editors for a given resource are closed, the `CustomDocument` is disposed of. Opening an editor at
 		 * this point will trigger another call to `resolveCustomDocument`.
 		 *
-		 * @param document Document to resolve.
+		 * @param uri Uri of the document to open.
 		 * @param token A cancellation token that indicates the result is no longer needed.
 		 *
-		 * @return The capabilities of the resolved document.
+		 * @return The custom document.
 		 */
-		resolveCustomDocument(document: CustomDocument, token: CancellationToken): Thenable<void>; // TODO: rename to open?
+		openCustomDocument(uri: Uri, token: CancellationToken): Thenable<CustomDocument<EditType>>;
 
 		/**
 		 * Resolve a webview editor for a given resource.
 		 *
-		 * This is called when a user first opens a resource for a `CustomTextEditorProvider`, or if they reopen an
+		 * This is called when a user first opens a resource for a `CustomEditorProvider`, or if they reopen an
-		 * existing editor using this `CustomTextEditorProvider`.
+		 * existing editor using this `CustomEditorProvider`.
 		 *
 		 * To resolve a webview editor, the provider must fill in its initial html content and hook up all
 		 * the event listeners it is interested it. The provider can also hold onto the `WebviewPanel` to use later,
@@ -1495,14 +1504,14 @@ declare module 'vscode' {
 		 *
 		 * @return Thenable indicating that the webview editor has been resolved.
 		 */
-		resolveCustomEditor(document: CustomDocument, webviewPanel: WebviewPanel, token: CancellationToken): Thenable<void>;
+		resolveCustomEditor(document: CustomDocument<EditType>, webviewPanel: WebviewPanel, token: CancellationToken): Thenable<void>;
 
 		/**
 		 * Defines the editing capability of a custom webview document.
 		 *
 		 * When not provided, the document is considered readonly.
 		 */
-		readonly editingDelegate?: CustomEditorEditingDelegate;
+		readonly editingDelegate?: CustomEditorEditingDelegate<EditType>;
 	}
 
 	/**
@@ -1516,6 +1525,7 @@ declare module 'vscode' {
 	 * For binary files or more specialized use cases, see [CustomEditorProvider](#CustomEditorProvider).
 	 */
 	export interface CustomTextEditorProvider {
+
 		/**
 		 * Resolve a webview editor for a given text resource.
 		 *
@@ -1549,8 +1559,6 @@ declare module 'vscode' {
 		 * @return Thenable indicating that the webview editor has been moved.
 		 */
 		moveCustomTextEditor?(newDocument: TextDocument, existingWebviewPanel: WebviewPanel, token: CancellationToken): Thenable<void>;
-
-		// TODO: handlesMove?: boolean;
 	}
 
 	namespace window {
@@ -1560,14 +1568,16 @@ declare module 'vscode' {
 		 * @param viewType Type of the webview editor provider. This should match the `viewType` from the
 		 *   `package.json` contributions.
 		 * @param provider Provider that resolves editors.
-		 * @param webviewOptions Content settings for the webview panels that the provider is given.
+		 * @param options Options for the provider
 		 *
 		 * @return Disposable that unregisters the provider.
 		 */
 		export function registerCustomEditorProvider(
 			viewType: string,
 			provider: CustomEditorProvider | CustomTextEditorProvider,
-			webviewOptions?: WebviewPanelOptions, // TODO: move this onto provider?
+			options?: {
+				readonly webviewOptions?: WebviewPanelOptions;
+			}
 		): Disposable;
 	}
 

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

@@ -663,13 +663,21 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		}
 
 		const undoneEdit = this._edits[this._currentEditIndex];
-		await this._proxy.$undo(this._realResource, this.viewType, undoneEdit);
+		await this._proxy.$undo(this._realResource, this.viewType, undoneEdit, this.getEditState());
 
 		this.change(() => {
 			--this._currentEditIndex;
 		});
 	}
 
+	private getEditState(): extHostProtocol.CustomDocumentEditState {
+		return {
+			allEdits: this._edits,
+			currentIndex: this._currentEditIndex,
+			saveIndex: this._savePoint,
+		};
+	}
+
 	private async redo(): Promise<void> {
 		if (!this._editable) {
 			return;
@@ -681,7 +689,7 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		}
 
 		const redoneEdit = this._edits[this._currentEditIndex + 1];
-		await this._proxy.$redo(this._realResource, this.viewType, redoneEdit);
+		await this._proxy.$redo(this._realResource, this.viewType, redoneEdit, this.getEditState());
 		this.change(() => {
 			++this._currentEditIndex;
 		});
@@ -728,7 +736,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._proxy.$revert(this._realResource, this.viewType, { undoneEdits: editsToUndo, redoneEdits: editsToRedo }, this.getEditState());
 		this.change(() => {
 			this._currentEditIndex = this._savePoint;
 			this.spliceEdits();

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

@@ -583,9 +583,9 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			registerWebviewPanelSerializer: (viewType: string, serializer: vscode.WebviewPanelSerializer) => {
 				return extHostWebviews.registerWebviewPanelSerializer(extension, viewType, serializer);
 			},
-			registerCustomEditorProvider: (viewType: string, provider: vscode.CustomEditorProvider | vscode.CustomTextEditorProvider, options?: vscode.WebviewPanelOptions) => {
+			registerCustomEditorProvider: (viewType: string, provider: vscode.CustomEditorProvider | vscode.CustomTextEditorProvider, options?: { webviewOptions?: vscode.WebviewPanelOptions }) => {
 				checkProposedApiEnabled(extension);
-				return extHostWebviews.registerCustomEditorProvider(extension, viewType, provider, options);
+				return extHostWebviews.registerCustomEditorProvider(extension, viewType, provider, options?.webviewOptions);
 			},
 			registerDecorationProvider(provider: vscode.DecorationProvider) {
 				checkProposedApiEnabled(extension);
@@ -1030,7 +1030,8 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			ColorThemeKind: extHostTypes.ColorThemeKind,
 			TimelineItem: extHostTypes.TimelineItem,
 			CellKind: extHostTypes.CellKind,
-			CellOutputKind: extHostTypes.CellOutputKind
+			CellOutputKind: extHostTypes.CellOutputKind,
+			CustomDocument: extHostTypes.CustomDocument,
 		};
 	};
 }

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

@@ -622,6 +622,12 @@ export interface WebviewPanelViewStateData {
 	};
 }
 
+export interface CustomDocumentEditState {
+	readonly allEdits: readonly number[];
+	readonly currentIndex: number;
+	readonly saveIndex: number;
+}
+
 export interface ExtHostWebviewsShape {
 	$onMessage(handle: WebviewPanelHandle, message: any): void;
 	$onMissingCsp(handle: WebviewPanelHandle, extensionId: string): void;
@@ -634,9 +640,9 @@ export interface ExtHostWebviewsShape {
 	$createWebviewCustomEditorDocument(resource: UriComponents, viewType: string, cancellation: CancellationToken): Promise<{ editable: boolean }>;
 	$disposeWebviewCustomEditorDocument(resource: UriComponents, viewType: string): Promise<void>;
 
-	$undo(resource: UriComponents, viewType: string, editId: number): Promise<void>;
+	$undo(resource: UriComponents, viewType: string, editId: number, state: CustomDocumentEditState): Promise<void>;
-	$redo(resource: UriComponents, viewType: string, editId: number): Promise<void>;
+	$redo(resource: UriComponents, viewType: string, editId: number, state: CustomDocumentEditState): Promise<void>;
-	$revert(resource: UriComponents, viewType: string, changes: { undoneEdits: number[], redoneEdits: number[] }): Promise<void>;
+	$revert(resource: UriComponents, viewType: string, changes: { undoneEdits: number[], redoneEdits: number[] }, state: CustomDocumentEditState): Promise<void>;
 	$disposeEdits(resourceComponents: UriComponents, viewType: string, editIds: number[]): void;
 
 	$onSave(resource: UriComponents, viewType: string, cancellation: CancellationToken): Promise<void>;

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

@@ -4,16 +4,19 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { coalesce, equals } from 'vs/base/common/arrays';
+import { escapeCodicons } from 'vs/base/common/codicons';
 import { illegalArgument } from 'vs/base/common/errors';
+import { Emitter } from 'vs/base/common/event';
 import { IRelativePattern } from 'vs/base/common/glob';
 import { isMarkdownString } from 'vs/base/common/htmlContent';
 import { startsWith } from 'vs/base/common/strings';
 import { URI } from 'vs/base/common/uri';
 import { generateUuid } from 'vs/base/common/uuid';
-import type * as vscode from 'vscode';
 import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/platform/files/common/files';
 import { RemoteAuthorityResolverErrorCode } from 'vs/platform/remote/common/remoteAuthorityResolver';
-import { escapeCodicons } from 'vs/base/common/codicons';
+import type * as vscode from 'vscode';
+import { Cache } from './cache';
+import { assertIsDefined } from 'vs/base/common/types';
 
 function es5ClassCompat(target: Function): any {
 	///@ts-ignore
@@ -2577,3 +2580,79 @@ export class TimelineItem implements vscode.TimelineItem {
 }
 
 //#endregion Timeline
+
+//#region Custom Editors
+
+interface EditState {
+	readonly allEdits: readonly number[];
+	readonly currentIndex: number;
+	readonly saveIndex: number;
+}
+
+export class CustomDocument<EditType = unknown> implements vscode.CustomDocument<EditType> {
+
+
+	readonly #edits = new Cache<EditType>('edits');
+
+	#editState: EditState;
+
+	readonly #viewType: string;
+	readonly #uri: vscode.Uri;
+
+	constructor(viewType: string, uri: vscode.Uri) {
+		this.#viewType = viewType;
+		this.#uri = uri;
+		this.#editState = {
+			allEdits: [],
+			currentIndex: 0,
+			saveIndex: 0
+		};
+	}
+
+	//#region Public API
+
+	public get viewType(): string { return this.#viewType; }
+
+	public get uri(): vscode.Uri { return this.#uri; }
+
+	#onDidDispose = new Emitter<void>();
+	public readonly onDidDispose = this.#onDidDispose.event;
+
+	get appliedEdits() {
+		return this.#editState.allEdits.slice(0, this.#editState.currentIndex)
+			.map(id => this._getEdit(id));
+	}
+
+	get savedEdits() {
+		return this.#editState.allEdits.slice(0, this.#editState.saveIndex)
+			.map(id => this._getEdit(id));
+	}
+
+	//#endregion
+
+	/** @internal */ _dispose(): void {
+		this.#onDidDispose.fire();
+		this.#onDidDispose.dispose();
+	}
+
+	/** @internal */ _updateEditState(state: EditState) {
+		this.#editState = state;
+	}
+
+	/** @internal*/ _getEdit(editId: number): EditType {
+		return assertIsDefined(this.#edits.get(editId, 0));
+	}
+
+	/** @internal*/ _disposeEdits(editIds: number[]) {
+		for (const editId of editIds) {
+			this.#edits.delete(editId);
+		}
+	}
+
+	/** @internal*/ _addEdit(edit: EditType): number {
+		return this.#edits.add([edit]);
+	}
+
+}
+
+// #endregion

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

@@ -18,9 +18,8 @@ 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 * as extHostProtocol from './extHost.protocol';
-import { ExtHostWebviewsShape, IMainContext, MainContext, MainThreadWebviewsShape, WebviewExtensionDescription, WebviewPanelHandle, WebviewPanelViewStateData } from './extHost.protocol';
+import * as extHostTypes from './extHostTypes';
-import { Disposable as VSCodeDisposable } from './extHostTypes';
 
 type IconPath = URI | { light: URI, dark: URI };
 
@@ -33,8 +32,8 @@ export class ExtHostWebview implements vscode.Webview {
 	public readonly onDidReceiveMessage: Event<any> = this._onMessageEmitter.event;
 
 	constructor(
-		private readonly _handle: WebviewPanelHandle,
+		private readonly _handle: extHostProtocol.WebviewPanelHandle,
-		private readonly _proxy: MainThreadWebviewsShape,
+		private readonly _proxy: extHostProtocol.MainThreadWebviewsShape,
 		private _options: vscode.WebviewOptions,
 		private readonly _initData: WebviewInitData,
 		private readonly _workspace: IExtHostWorkspace | undefined,
@@ -99,8 +98,8 @@ export class ExtHostWebview implements vscode.Webview {
 
 export class ExtHostWebviewEditor extends Disposable implements vscode.WebviewPanel {
 
-	private readonly _handle: WebviewPanelHandle;
+	private readonly _handle: extHostProtocol.WebviewPanelHandle;
-	private readonly _proxy: MainThreadWebviewsShape;
+	private readonly _proxy: extHostProtocol.MainThreadWebviewsShape;
 	private readonly _viewType: string;
 	private _title: string;
 	private _iconPath?: IconPath;
@@ -121,8 +120,8 @@ export class ExtHostWebviewEditor extends Disposable implements vscode.WebviewPa
 	public readonly onDidChangeViewState = this.#onDidChangeViewState.event;
 
 	constructor(
-		handle: WebviewPanelHandle,
+		handle: extHostProtocol.WebviewPanelHandle,
-		proxy: MainThreadWebviewsShape,
+		proxy: extHostProtocol.MainThreadWebviewsShape,
 		viewType: string,
 		title: string,
 		viewColumn: vscode.ViewColumn | undefined,
@@ -246,114 +245,14 @@ export class ExtHostWebviewEditor extends Disposable implements vscode.WebviewPa
 	}
 }
 
-class CustomDocument extends Disposable implements vscode.CustomDocument {
-
-	public static create(
-		viewType: string,
-		uri: vscode.Uri,
-		editingDelegate: vscode.CustomEditorEditingDelegate | undefined
-	) {
-		return Object.seal(new CustomDocument(viewType, uri, editingDelegate));
-	}
-
-	// Explicitly initialize all properties as we seal the object after creation!
-
-	readonly #_edits = new Cache<unknown>('edits');
-
-	readonly #viewType: string;
-	readonly #uri: vscode.Uri;
-	readonly #editingDelegate: vscode.CustomEditorEditingDelegate | undefined;
-
-	private constructor(
-		viewType: string,
-		uri: vscode.Uri,
-		editingDelegate: vscode.CustomEditorEditingDelegate | undefined,
-	) {
-		super();
-		this.#viewType = viewType;
-		this.#uri = uri;
-		this.#editingDelegate = editingDelegate;
-	}
-
-	dispose() {
-		this.#onDidDispose.fire();
-		super.dispose();
-	}
-
-	//#region Public API
-
-	public get viewType(): string { return this.#viewType; }
-
-	public get uri(): vscode.Uri { return this.#uri; }
-
-	#onDidDispose = this._register(new Emitter<void>());
-	public readonly onDidDispose = this.#onDidDispose.event;
-
-	public userData: unknown = undefined;
-
-	//#endregion
-
-	//#region Internal
-
-	/** @internal*/ async _revert(changes: { undoneEdits: number[], redoneEdits: number[] }) {
-		const editing = this.getEditingDelegate();
-		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(this, { undoneEdits, appliedEdits });
-	}
-
-	/** @internal*/ _undo(editId: number) {
-		const editing = this.getEditingDelegate();
-		const edit = this.#_edits.get(editId, 0);
-		return editing.undoEdits(this, [edit]);
-	}
-
-	/** @internal*/ _redo(editId: number) {
-		const editing = this.getEditingDelegate();
-		const edit = this.#_edits.get(editId, 0);
-		return editing.applyEdits(this, [edit]);
-	}
-
-	/** @internal*/ _save(cancellation: CancellationToken) {
-		return this.getEditingDelegate().save(this, cancellation);
-	}
-
-	/** @internal*/ _saveAs(target: vscode.Uri) {
-		return this.getEditingDelegate().saveAs(this, target);
-	}
-
-	/** @internal*/ _backup(cancellation: CancellationToken) {
-		return this.getEditingDelegate().backup(this, cancellation);
-	}
-
-	/** @internal*/ _disposeEdits(editIds: number[]) {
-		for (const editId of editIds) {
-			this.#_edits.delete(editId);
-		}
-	}
-
-	/** @internal*/ _pushEdit(edit: unknown): number {
-		return this.#_edits.add([edit]);
-	}
-
-	//#endregion
-
-	private getEditingDelegate(): vscode.CustomEditorEditingDelegate {
-		if (!this.#editingDelegate) {
-			throw new Error('Document is not editable');
-		}
-		return this.#editingDelegate;
-	}
-}
-
 class WebviewDocumentStore {
-	private readonly _documents = new Map<string, CustomDocument>();
+	private readonly _documents = new Map<string, extHostTypes.CustomDocument>();
 
-	public get(viewType: string, resource: vscode.Uri): CustomDocument | undefined {
+	public get(viewType: string, resource: vscode.Uri): extHostTypes.CustomDocument | undefined {
 		return this._documents.get(this.key(viewType, resource));
 	}
 
-	public add(document: CustomDocument) {
+	public add(document: extHostTypes.CustomDocument) {
 		const key = this.key(document.viewType, document.uri);
 		if (this._documents.has(key)) {
 			throw new Error(`Document already exists for viewType:${document.viewType} resource:${document.uri}`);
@@ -361,7 +260,7 @@ class WebviewDocumentStore {
 		this._documents.set(key, document);
 	}
 
-	public delete(document: CustomDocument) {
+	public delete(document: extHostTypes.CustomDocument) {
 		const key = this.key(document.viewType, document.uri);
 		this._documents.delete(key);
 	}
@@ -406,18 +305,18 @@ class EditorProviderStore {
 			throw new Error(`Provider for viewType:${viewType} already registered`);
 		}
 		this._providers.set(viewType, { type, extension, provider } as ProviderEntry);
-		return new VSCodeDisposable(() => this._providers.delete(viewType));
+		return new extHostTypes.Disposable(() => this._providers.delete(viewType));
 	}
 }
 
-export class ExtHostWebviews implements ExtHostWebviewsShape {
+export class ExtHostWebviews implements extHostProtocol.ExtHostWebviewsShape {
 
-	private static newHandle(): WebviewPanelHandle {
+	private static newHandle(): extHostProtocol.WebviewPanelHandle {
 		return generateUuid();
 	}
 
-	private readonly _proxy: MainThreadWebviewsShape;
+	private readonly _proxy: extHostProtocol.MainThreadWebviewsShape;
-	private readonly _webviewPanels = new Map<WebviewPanelHandle, ExtHostWebviewEditor>();
+	private readonly _webviewPanels = new Map<extHostProtocol.WebviewPanelHandle, ExtHostWebviewEditor>();
 
 	private readonly _serializers = new Map<string, {
 		readonly serializer: vscode.WebviewPanelSerializer;
@@ -429,14 +328,14 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	private readonly _documents = new WebviewDocumentStore();
 
 	constructor(
-		mainContext: IMainContext,
+		mainContext: extHostProtocol.IMainContext,
 		private readonly initData: WebviewInitData,
 		private readonly workspace: IExtHostWorkspace | undefined,
 		private readonly _logService: ILogService,
 		private readonly _deprecationService: IExtHostApiDeprecationService,
 		private readonly _extHostDocuments: ExtHostDocuments,
 	) {
-		this._proxy = mainContext.getProxy(MainContext.MainThreadWebviews);
+		this._proxy = mainContext.getProxy(extHostProtocol.MainContext.MainThreadWebviews);
 	}
 
 	public createWebviewPanel(
@@ -473,7 +372,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		this._serializers.set(viewType, { serializer, extension });
 		this._proxy.$registerSerializer(viewType);
 
-		return new VSCodeDisposable(() => {
+		return new extHostTypes.Disposable(() => {
 			this._serializers.delete(viewType);
 			this._proxy.$unregisterSerializer(viewType);
 		});
@@ -497,21 +396,21 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 			if (provider.editingDelegate) {
 				disposables.add(provider.editingDelegate.onDidEdit(e => {
 					const document = e.document;
-					const editId = (document as CustomDocument)._pushEdit(e.edit);
+					const editId = (document as unknown as extHostTypes.CustomDocument)._addEdit(e.edit);
 					this._proxy.$onDidEdit(document.uri, document.viewType, editId, e.label);
 				}));
 			}
 		}
 
-		return VSCodeDisposable.from(
+		return extHostTypes.Disposable.from(
 			disposables,
-			new VSCodeDisposable(() => {
+			new extHostTypes.Disposable(() => {
 				this._proxy.$unregisterEditorProvider(viewType);
 			}));
 	}
 
 	public $onMessage(
-		handle: WebviewPanelHandle,
+		handle: extHostProtocol.WebviewPanelHandle,
 		message: any
 	): void {
 		const panel = this.getWebviewPanel(handle);
@@ -521,13 +420,13 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	}
 
 	public $onMissingCsp(
-		_handle: WebviewPanelHandle,
+		_handle: extHostProtocol.WebviewPanelHandle,
 		extensionId: string
 	): void {
 		this._logService.warn(`${extensionId} created a webview without a content security policy: https://aka.ms/vscode-webview-missing-csp`);
 	}
 
-	public $onDidChangeWebviewPanelViewStates(newStates: WebviewPanelViewStateData): void {
+	public $onDidChangeWebviewPanelViewStates(newStates: extHostProtocol.WebviewPanelViewStateData): void {
 		const handles = Object.keys(newStates);
 		// Notify webviews of state changes in the following order:
 		// - Non-visible
@@ -560,7 +459,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		}
 	}
 
-	async $onDidDisposeWebviewPanel(handle: WebviewPanelHandle): Promise<void> {
+	async $onDidDisposeWebviewPanel(handle: extHostProtocol.WebviewPanelHandle): Promise<void> {
 		const panel = this.getWebviewPanel(handle);
 		if (panel) {
 			panel.dispose();
@@ -569,7 +468,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	}
 
 	async $deserializeWebviewPanel(
-		webviewHandle: WebviewPanelHandle,
+		webviewHandle: extHostProtocol.WebviewPanelHandle,
 		viewType: string,
 		title: string,
 		state: any,
@@ -599,9 +498,8 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		}
 
 		const revivedResource = URI.revive(resource);
-		const document = CustomDocument.create(viewType, revivedResource, entry.provider.editingDelegate);
+		const document = await entry.provider.openCustomDocument(revivedResource, cancellation);
-		await entry.provider.resolveCustomDocument(document, cancellation);
+		this._documents.add(document as unknown as extHostTypes.CustomDocument);
-		this._documents.add(document);
 		return {
 			editable: !!entry.provider.editingDelegate,
 		};
@@ -620,12 +518,12 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		const revivedResource = URI.revive(resource);
 		const document = this.getCustomDocument(viewType, revivedResource);
 		this._documents.delete(document);
-		document.dispose();
+		document._dispose();
 	}
 
 	async $resolveWebviewEditor(
 		resource: UriComponents,
-		handle: WebviewPanelHandle,
+		handle: extHostProtocol.WebviewPanelHandle,
 		viewType: string,
 		title: string,
 		position: EditorViewColumn,
@@ -686,50 +584,71 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		await (entry.provider as vscode.CustomTextEditorProvider).moveCustomTextEditor!(document, webview, CancellationToken.None);
 	}
 
-	async $undo(resourceComponents: UriComponents, viewType: string, editId: number): Promise<void> {
+	async $undo(resourceComponents: UriComponents, viewType: string, editId: number, state: extHostProtocol.CustomDocumentEditState): Promise<void> {
+		const delegate = this.getEditingDelegate(viewType);
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		return document._undo(editId);
+		document._updateEditState(state);
+		return delegate.undoEdits(document, [document._getEdit(editId)]);
 	}
 
-	async $redo(resourceComponents: UriComponents, viewType: string, editId: number): Promise<void> {
+	async $redo(resourceComponents: UriComponents, viewType: string, editId: number, state: extHostProtocol.CustomDocumentEditState): Promise<void> {
+		const delegate = this.getEditingDelegate(viewType);
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		return document._redo(editId);
+		return delegate.applyEdits(document, [document._getEdit(editId)]);
 	}
 
 	async $revert(resourceComponents: UriComponents, viewType: string, changes: { undoneEdits: number[], redoneEdits: number[] }): Promise<void> {
+		const delegate = this.getEditingDelegate(viewType);
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		return document._revert(changes);
+		const undoneEdits = changes.undoneEdits.map(id => document._getEdit(id));
+		const appliedEdits = changes.redoneEdits.map(id => document._getEdit(id));
+		return delegate.revert(document, { undoneEdits, appliedEdits });
 	}
 
 	async $onSave(resourceComponents: UriComponents, viewType: string, cancellation: CancellationToken): Promise<void> {
+		const delegate = this.getEditingDelegate(viewType);
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		return document._save(cancellation);
+		return delegate.save(document, cancellation);
 	}
 
 	async $onSaveAs(resourceComponents: UriComponents, viewType: string, targetResource: UriComponents): Promise<void> {
+		const delegate = this.getEditingDelegate(viewType);
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		return document._saveAs(URI.revive(targetResource));
+		return delegate.saveAs(document, URI.revive(targetResource));
 	}
 
 	async $backup(resourceComponents: UriComponents, viewType: string, cancellation: CancellationToken): Promise<void> {
+		const delegate = this.getEditingDelegate(viewType);
 		const document = this.getCustomDocument(viewType, resourceComponents);
-		return document._backup(cancellation);
+		return delegate.backup(document, cancellation);
 	}
 
-	private getWebviewPanel(handle: WebviewPanelHandle): ExtHostWebviewEditor | undefined {
+	private getWebviewPanel(handle: extHostProtocol.WebviewPanelHandle): ExtHostWebviewEditor | undefined {
 		return this._webviewPanels.get(handle);
 	}
 
-	private getCustomDocument(viewType: string, resource: UriComponents): CustomDocument {
+	private getCustomDocument(viewType: string, resource: UriComponents): extHostTypes.CustomDocument {
 		const document = this._documents.get(viewType, URI.revive(resource));
 		if (!document) {
 			throw new Error('No webview editor custom document found');
 		}
 		return document;
 	}
+
+	private getEditingDelegate(viewType: string): vscode.CustomEditorEditingDelegate {
+		const entry = this._editorProviders.get(viewType);
+		if (!entry) {
+			throw new Error(`No provider found for '${viewType}'`);
+		}
+		const delegate = (entry.provider as vscode.CustomEditorProvider).editingDelegate;
+		if (!delegate) {
+			throw new Error(`Provider for ${viewType}' does not support editing`);
+		}
+		return delegate;
+	}
 }
 
-function toExtensionData(extension: IExtensionDescription): WebviewExtensionDescription {
+function toExtensionData(extension: IExtensionDescription): extHostProtocol.WebviewExtensionDescription {
 	return { id: extension.identifier, location: extension.extensionLocation };
 }