Migrate webview notebook editor backup to builtin serializer (#129211)

* Migrate webview notebook editor backup to builtin serializer

* update extension kind.

* discard webview backup once it is opened in notebook editor.

extensions/ipynb/package.json

@@ -53,6 +53,7 @@
     "watch": "npx gulp watch-extension:ipynb"
   },
 	"dependencies": {
+		"@enonic/fnv-plus": "^1.3.0",
 		"detect-indent": "^6.0.0"
 	},
 	"devDependencies": {

extensions/ipynb/src/ipynbMain.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode';
 import { NotebookSerializer } from './serializer';
 
 export function activate(context: vscode.ExtensionContext) {
-	context.subscriptions.push(vscode.workspace.registerNotebookSerializer('jupyter-notebook', new NotebookSerializer(), {
+	context.subscriptions.push(vscode.workspace.registerNotebookSerializer('jupyter-notebook', new NotebookSerializer(context), {
 		transientOutputs: false,
 		transientCellMetadata: {
 			breakpointMargin: true,

extensions/ipynb/src/serializer.ts

@@ -8,9 +8,13 @@ import * as detectIndent from 'detect-indent';
 import * as vscode from 'vscode';
 import { defaultNotebookFormat } from './constants';
 import { createJupyterCellFromNotebookCell, getPreferredLanguage, jupyterNotebookModelToNotebookData, pruneCell } from './helpers';
+import * as fnv from '@enonic/fnv-plus';
 
 export class NotebookSerializer implements vscode.NotebookSerializer {
-	public deserializeNotebook(content: Uint8Array, _token: vscode.CancellationToken): vscode.NotebookData {
+	constructor(readonly context: vscode.ExtensionContext) {
+	}
+
+	public async deserializeNotebook(content: Uint8Array, _token: vscode.CancellationToken): Promise<vscode.NotebookData> {
 		let contents = '';
 		try {
 			contents = new TextDecoder().decode(content);
@@ -20,6 +24,22 @@ export class NotebookSerializer implements vscode.NotebookSerializer {
 		let json: Partial<nbformat.INotebookContent>;
 		try {
 			json = contents ? (JSON.parse(contents) as Partial<nbformat.INotebookContent>) : {};
+
+			if (json.__webview_backup) {
+				const backupId = json.__webview_backup;
+				const uri = this.context.globalStorageUri;
+				const folder = uri.with({ path: this.context.globalStorageUri.path.replace('vscode.ipynb', 'ms-toolsai.jupyter') });
+				const fileHash = fnv.fast1a32hex(backupId) as string;
+				const fileName = `${fileHash}.ipynb`;
+				const file = vscode.Uri.joinPath(folder, fileName);
+				const data = await vscode.workspace.fs.readFile(file);
+				json = data ? JSON.parse(data.toString()) : {};
+
+				if (json.contents && typeof json.contents === 'string') {
+					contents = json.contents;
+					json = JSON.parse(contents) as Partial<nbformat.INotebookContent>;
+				}
+			}
 		} catch (e) {
 			console.log(contents);
 			console.log(e);

extensions/ipynb/src/types.d.ts

@@ -6,3 +6,5 @@
 
 /// <reference path='../../../src/vs/vscode.d.ts'/>
 /// <reference path='../../../src/vs/vscode.proposed.d.ts'/>
+
+declare module '@enonic/fnv-plus';

extensions/ipynb/yarn.lock

@@ -2,6 +2,11 @@
 # yarn lockfile v1
 
 
+"@enonic/fnv-plus@^1.3.0":
+  version "1.3.0"
+  resolved "https://registry.yarnpkg.com/@enonic/fnv-plus/-/fnv-plus-1.3.0.tgz#be65a7b128a3b544f60aea3ef978d938e85869f3"
+  integrity sha512-BCN9uNWH8AmiP7BXBJqEinUY9KXalmRzo+L0cB/mQsmFfzODxwQrbvxCHXUNH2iP+qKkWYtB4vyy8N62PViMFw==
+
 "@jupyterlab/coreutils@^3.1.0":
   version "3.2.0"
   resolved "https://registry.yarnpkg.com/@jupyterlab/coreutils/-/coreutils-3.2.0.tgz#dd4d887bdedfea4c8545d46d297531749cb13724"

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

@@ -17,6 +17,7 @@ import { IWorkingCopyEditorService } from 'vs/workbench/services/workingCopy/com
 import { ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
 import { Schemas } from 'vs/base/common/network';
 import { isEqual } from 'vs/base/common/resources';
+import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebookEditorInput';
 
 export interface CustomDocumentBackupData extends IWorkingCopyBackupMeta {
 	readonly viewType: string;
@@ -88,6 +89,10 @@ export class CustomEditorInputSerializer extends WebviewEditorInputSerializer {
 		serializedEditorInput: string
 	): CustomEditorInput {
 		const data = this.fromJson(JSON.parse(serializedEditorInput));
+		if (data.viewType === 'jupyter.notebook.ipynb') {
+			return NotebookEditorInput.create(this._instantiationService, data.editorResource, 'jupyter-notebook', { _backupId: data.backupId }) as any;
+		}
+
 		const webview = reviveWebview(this._webviewService, data);
 		const customInput = this._instantiationService.createInstance(CustomEditorInput, data.editorResource, data.viewType, data.id, webview, { startsDirty: data.dirty, backupId: data.backupId });
 		if (typeof data.group === 'number') {
@@ -125,6 +130,15 @@ export class ComplexCustomWorkingCopyEditorHandler extends Disposable implements
 		this._register(this._workingCopyEditorService.registerHandler({
 			handles: workingCopy => workingCopy.resource.scheme === Schemas.vscodeCustomEditor,
 			isOpen: (workingCopy, editor) => {
+				if (workingCopy.resource.authority === 'jupyter-notebook-ipynb' && editor instanceof NotebookEditorInput) {
+					try {
+						const data = JSON.parse(workingCopy.resource.query);
+						const workingCopyResource = URI.from(data);
+						return isEqual(workingCopyResource, editor.resource);
+					} catch {
+						return false;
+					}
+				}
 				if (!(editor instanceof CustomEditorInput)) {
 					return false;
 				}
@@ -149,6 +163,10 @@ export class ComplexCustomWorkingCopyEditorHandler extends Disposable implements
 				}
 
 				const backupData = backup.meta;
+				if (backupData.viewType === 'jupyter.notebook.ipynb') {
+					return NotebookEditorInput.create(this._instantiationService, URI.revive(backupData.editorResource), 'jupyter-notebook', { _backupId: backupData.backupId, _workingCopy: workingCopy }) as any;
+				}
+
 				const id = backupData.webview.id;
 				const extension = reviveWebviewExtensionDescription(backupData.extension?.id, backupData.extension?.location);
 				const webview = reviveWebview(this._webviewService, {

src/vs/workbench/contrib/notebook/common/model/notebookTextModel.ts

@@ -285,7 +285,7 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
 		);
 	}
 
-	private _initialize(cells: ICellDto2[]) {
+	_initialize(cells: ICellDto2[], triggerDirty?: boolean) {
 		this._cells = [];
 		this._versionId = 0;
 		this._notebookSpecificAlternativeId = 0;
@@ -306,6 +306,10 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
 
 		this._cells.splice(0, 0, ...mainCells);
 		this._alternativeVersionId = this._generateAlternativeId();
+
+		if (triggerDirty) {
+			this._eventEmitter.emit({ kind: NotebookCellsChangeType.Unknown, transient: false }, true);
+		}
 	}
 
 	private _bindCellContentHandler(cell: NotebookCellTextModel, e: 'content' | 'language' | 'mime') {

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

@@ -5,7 +5,7 @@
 
 import * as glob from 'vs/base/common/glob';
 import { IEditorInput, GroupIdentifier, ISaveOptions, IMoveResult, IRevertOptions, EditorInputCapabilities, Verbosity, IUntypedEditorInput } from 'vs/workbench/common/editor';
-import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
+import { INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { URI } from 'vs/base/common/uri';
 import { isEqual, joinPath } from 'vs/base/common/resources';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
@@ -20,9 +20,17 @@ import { FileSystemProviderCapabilities, IFileService } from 'vs/platform/files/
 import { AbstractResourceEditorInput } from 'vs/workbench/common/editor/resourceEditorInput';
 import { IResourceEditorInput } from 'vs/platform/editor/common/editor';
 import { onUnexpectedError } from 'vs/base/common/errors';
+import { VSBuffer } from 'vs/base/common/buffer';
+import { IWorkingCopyIdentifier } from 'vs/workbench/services/workingCopy/common/workingCopy';
+import { IWorkingCopyBackupService } from 'vs/workbench/services/workingCopy/common/workingCopyBackup';
 
 interface NotebookEditorInputOptions {
 	startDirty?: boolean;
+	/**
+	 * backupId for webview
+	 */
+	_backupId?: string;
+	_workingCopy?: IWorkingCopyIdentifier;
 }
 
 export class NotebookEditorInput extends AbstractResourceEditorInput {
@@ -45,6 +53,7 @@ export class NotebookEditorInput extends AbstractResourceEditorInput {
 		@INotebookEditorModelResolverService private readonly _notebookModelResolverService: INotebookEditorModelResolverService,
 		@IFileDialogService private readonly _fileDialogService: IFileDialogService,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
+		@IWorkingCopyBackupService private readonly workingCopyBackupService: IWorkingCopyBackupService,
 		@ILabelService labelService: ILabelService,
 		@IFileService fileService: IFileService
 	) {
@@ -232,6 +241,20 @@ export class NotebookEditorInput extends AbstractResourceEditorInput {
 			this._editorModelReference.object.load();
 		}
 
+		if (this.options._backupId) {
+			const info = await this._notebookService.withNotebookDataProvider(this._editorModelReference.object.notebook.uri, this._editorModelReference.object.notebook.viewType);
+			if (!(info instanceof SimpleNotebookProviderInfo)) {
+				throw new Error('CANNOT open file notebook with this provider');
+			}
+
+			const data = await info.serializer.dataToNotebook(VSBuffer.fromString(JSON.stringify({ __webview_backup: this.options._backupId })));
+			this._editorModelReference.object.notebook._initialize(data.cells, true);
+
+			if (this.options._workingCopy) {
+				await this.workingCopyBackupService.discardBackup(this.options._workingCopy);
+			}
+		}
+
 		return this._editorModelReference.object;
 	}