Merge pull request #102845 from rebornix/rebornix/shared-notebook-models

shared notebook models

src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts

@@ -10,7 +10,7 @@ import { notebookProviderExtensionPoint, notebookRendererExtensionPoint, INotebo
 import { NotebookProviderInfo, NotebookEditorDescriptor } from 'vs/workbench/contrib/notebook/common/notebookProvider';
 import { NotebookExtensionDescription } from 'vs/workbench/api/common/extHost.protocol';
 import { Emitter, Event } from 'vs/base/common/event';
-import { INotebookTextModel, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, CellOutputKind, ITransformedDisplayOutputDto, IDisplayOutput, ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, NOTEBOOK_DISPLAY_ORDER, sortMimeTypes, IOrderedMimeType, mimeTypeSupportedByCore, IOutputRenderRequestOutputInfo, IOutputRenderRequestCellInfo, NotebookCellOutputsSplice, ICellEditOperation, CellEditType, ICellInsertEdit, IOutputRenderResponse, IProcessedOutput, BUILTIN_RENDERER_ID, NotebookEditorPriority, INotebookKernelProvider, notebookDocumentFilterMatch, INotebookKernelInfo2 } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookRendererInfo, INotebookKernelInfo, CellOutputKind, ITransformedDisplayOutputDto, IDisplayOutput, ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, NOTEBOOK_DISPLAY_ORDER, sortMimeTypes, IOrderedMimeType, mimeTypeSupportedByCore, IOutputRenderRequestOutputInfo, IOutputRenderRequestCellInfo, NotebookCellOutputsSplice, ICellEditOperation, CellEditType, ICellInsertEdit, IOutputRenderResponse, IProcessedOutput, BUILTIN_RENDERER_ID, NotebookEditorPriority, INotebookKernelProvider, notebookDocumentFilterMatch, INotebookKernelInfo2 } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 import { NotebookOutputRendererInfo } from 'vs/workbench/contrib/notebook/common/notebookOutputRenderer';
 import { Iterable } from 'vs/base/common/iterator';
@@ -29,6 +29,7 @@ import { StorageScope, IStorageService } from 'vs/platform/storage/common/storag
 import { IExtensionPointUser } from 'vs/workbench/services/extensions/common/extensionsRegistry';
 import { generateUuid } from 'vs/base/common/uuid';
 import { flatten } from 'vs/base/common/arrays';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 
 function MODEL_ID(resource: URI): string {
 	return resource.toString();
@@ -166,6 +167,7 @@ class ModelData implements IDisposable {
 }
 export class NotebookService extends Disposable implements INotebookService, ICustomEditorViewTypesHandler {
 	declare readonly _serviceBrand: undefined;
+	static mainthreadNotebookDocumentHandle: number = 0;
 	private readonly _notebookProviders = new Map<string, { controller: IMainNotebookController, extensionData: NotebookExtensionDescription }>();
 	private readonly _notebookRenderers = new Map<string, INotebookRendererInfo>();
 	private readonly _notebookKernels = new Map<string, INotebookKernelInfo>();
@@ -205,7 +207,8 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		@IEditorService private readonly _editorService: IEditorService,
 		@IConfigurationService private readonly _configurationService: IConfigurationService,
 		@IAccessibilityService private readonly _accessibilityService: IAccessibilityService,
-		@IStorageService private readonly _storageService: IStorageService
+		@IStorageService private readonly _storageService: IStorageService,
+		@IInstantiationService private readonly _instantiationService: IInstantiationService
 	) {
 		super();
 
@@ -393,43 +396,33 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		return renderer;
 	}
 
-	async createNotebookFromBackup(viewType: string, uri: URI, metadata: NotebookDocumentMetadata, languages: string[], cells: ICellDto2[], editorId?: string): Promise<NotebookTextModel | undefined> {
+	async resolveNotebook(viewType: string, uri: URI, forceReload: boolean, editorId?: string, backupId?: string): Promise<NotebookTextModel | undefined> {
 		const provider = this._notebookProviders.get(viewType);
 		if (!provider) {
 			return undefined;
 		}
 
-		const notebookModel = await provider.controller.createNotebook(viewType, uri, { metadata, languages, cells }, false, editorId);
-		if (!notebookModel) {
-			return undefined;
-		}
-
-		// new notebook model created
 		const modelId = MODEL_ID(uri);
-		const modelData = new ModelData(
-			notebookModel,
-			(model) => this._onWillDisposeDocument(model),
-		);
-		this._models.set(modelId, modelData);
-		this._onNotebookDocumentAdd.fire([notebookModel.uri]);
-		// after the document is added to the store and sent to ext host, we transform the ouputs
-		await this.transformTextModelOutputs(notebookModel!);
-		return modelData.model;
-	}
 
-	async resolveNotebook(viewType: string, uri: URI, forceReload: boolean, editorId?: string, backupId?: string): Promise<NotebookTextModel | undefined> {
-		const provider = this._notebookProviders.get(viewType);
-		if (!provider) {
-			return undefined;
-		}
+		let notebookModel: NotebookTextModel | undefined = undefined;
+		if (this._models.has(modelId)) {
+			// the model already exists
+			notebookModel = this._models.get(modelId)!.model;
+			if (forceReload) {
+				await provider.controller.reloadNotebook(notebookModel);
+			}
 
-		const notebookModel = await provider.controller.createNotebook(viewType, uri, undefined, forceReload, editorId, backupId);
-		if (!notebookModel) {
-			return undefined;
+			return notebookModel;
+		} else {
+			notebookModel = this._instantiationService.createInstance(NotebookTextModel, NotebookService.mainthreadNotebookDocumentHandle++, viewType, provider.controller.supportBackup, uri);
+			await provider.controller.createNotebook(notebookModel, backupId);
+
+			if (!notebookModel) {
+				return undefined;
+			}
 		}
 
 		// new notebook model created
-		const modelId = MODEL_ID(uri);
 		const modelData = new ModelData(
 			notebookModel!,
 			(model) => this._onWillDisposeDocument(model),
@@ -447,6 +440,12 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		return modelData.model;
 	}
 
+	getNotebookTextModel(uri: URI): NotebookTextModel | undefined {
+		const modelId = MODEL_ID(uri);
+
+		return this._models.get(modelId)?.model;
+	}
+
 	private async _fillInTransformedOutputs<T>(
 		renderers: Set<string>,
 		requestItems: IOutputRenderRequestCellInfo<T>[],
@@ -861,7 +860,8 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 			let provider = this._notebookProviders.get(modelData!.model.viewType);
 
 			if (provider) {
-				provider.controller.removeNotebookDocument(modelData!.model);
+				provider.controller.removeNotebookDocument(modelData!.model.uri);
+				modelData!.model.dispose();
 			}
 
 

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

@@ -3,13 +3,14 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import * as nls from 'vs/nls';
 import { Emitter, Event } from 'vs/base/common/event';
 import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { INotebookTextModel, NotebookCellOutputsSplice, NotebookCellTextModelSplice, NotebookDocumentMetadata, NotebookCellMetadata, ICellEditOperation, CellEditType, CellUri, ICellInsertEdit, NotebookCellsChangedEvent, CellKind, IProcessedOutput, notebookDocumentMetadataDefaults, diff, ICellDeleteEdit, NotebookCellsChangeType, ICellDto2, IMainCellDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { ITextSnapshot } from 'vs/editor/common/model';
-import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
+import { IUndoRedoService, UndoRedoElementType } from 'vs/platform/undoRedo/common/undoRedo';
 import { InsertCellEdit, DeleteCellEdit, MoveCellEdit, SpliceCellsEdit } from 'vs/workbench/contrib/notebook/common/model/cellEdit';
 import { ITextModelService } from 'vs/editor/common/services/resolverService';
 
@@ -116,8 +117,8 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
 		public viewType: string,
 		public supportBackup: boolean,
 		public uri: URI,
-		private _undoService: IUndoRedoService,
-		private _modelService: ITextModelService
+		@IUndoRedoService private _undoService: IUndoRedoService,
+		@ITextModelService private _modelService: ITextModelService
 	) {
 		super();
 		this.cells = [];
@@ -172,7 +173,7 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
 		this._increaseVersionId();
 	}
 
-	$applyEdit(modelVersionId: number, rawEdits: ICellEditOperation[], emitToExtHost: boolean, synchronous: boolean): boolean {
+	$applyEdit(modelVersionId: number, rawEdits: ICellEditOperation[], synchronous: boolean): boolean {
 		if (modelVersionId !== this._versionId) {
 			return false;
 		}
@@ -233,22 +234,20 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
 			return [diff.start, diff.deleteCount, diff.toInsert] as [number, number, NotebookCellTextModel[]];
 		});
 
-		if (emitToExtHost) {
-			this._onDidModelChangeProxy.fire({
-				kind: NotebookCellsChangeType.ModelChange,
-				versionId: this._versionId,
-				changes: diffs.map(diff => [diff[0], diff[1], diff[2].map(cell => ({
-					handle: cell.handle,
-					uri: cell.uri,
-					source: cell.textBuffer.getLinesContent(),
-					eol: cell.textBuffer.getEOL(),
-					language: cell.language,
-					cellKind: cell.cellKind,
-					outputs: cell.outputs,
-					metadata: cell.metadata
-				}))] as [number, number, IMainCellDto[]])
-			});
-		}
+		this._onDidModelChangeProxy.fire({
+			kind: NotebookCellsChangeType.ModelChange,
+			versionId: this._versionId,
+			changes: diffs.map(diff => [diff[0], diff[1], diff[2].map(cell => ({
+				handle: cell.handle,
+				uri: cell.uri,
+				source: cell.textBuffer.getLinesContent(),
+				eol: cell.textBuffer.getEOL(),
+				language: cell.language,
+				cellKind: cell.cellKind,
+				outputs: cell.outputs,
+				metadata: cell.metadata
+			}))] as [number, number, IMainCellDto[]])
+		});
 
 		const undoDiff = diffs.map(diff => {
 			const deletedCells = this.cells.slice(diff[0], diff[0] + diff[1]);
@@ -266,6 +265,21 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
 		return true;
 	}
 
+	$handleEdit(label: string | undefined, undo: () => void, redo: () => void): void {
+		this._undoService.pushElement({
+			type: UndoRedoElementType.Resource,
+			resource: this.uri,
+			label: label ?? nls.localize('defaultEditLabel', "Edit"),
+			undo: async () => {
+				undo();
+			},
+			redo: async () => {
+				redo();
+			},
+		});
+		this.setDirty(true);
+	}
+
 	createSnapshot(preserveBOM?: boolean): ITextSnapshot {
 		return new NotebookTextModelSnapshot(this);
 	}

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

@@ -14,7 +14,6 @@ import { IWorkingCopyService, IWorkingCopy, IWorkingCopyBackup, WorkingCopyCapab
 import { basename } from 'vs/base/common/resources';
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
 import { IBackupFileService } from 'vs/workbench/services/backup/common/backup';
-import { DefaultEndOfLine, ITextBuffer, EndOfLinePreference } from 'vs/editor/common/model';
 import { Schemas } from 'vs/base/common/network';
 import { IFileStatWithMetadata, IFileService } from 'vs/platform/files/common/files';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
@@ -142,42 +141,9 @@ export class NotebookEditorModel extends EditorModel implements IWorkingCopy, IN
 			return this; // Make sure meanwhile someone else did not succeed in loading
 		}
 
-		if (backup && backup.meta?.backupId === undefined) {
-			try {
-				return await this.loadFromBackup(backup.value.create(DefaultEndOfLine.LF), options?.editorId);
-			} catch (error) {
-				// this.logService.error('[text file model] load() from backup', error); // ignore error and continue to load as file below
-			}
-		}
-
 		return this.loadFromProvider(false, options?.editorId, backup?.meta?.backupId);
 	}
 
-	private async loadFromBackup(content: ITextBuffer, editorId?: string): Promise<NotebookEditorModel> {
-		const fullRange = content.getRangeAt(0, content.getLength());
-		const data = JSON.parse(content.getValueInRange(fullRange, EndOfLinePreference.LF));
-
-		const notebook = await this._notebookService.createNotebookFromBackup(this.viewType!, this.resource, data.metadata, data.languages, data.cells, editorId);
-		this._notebook = notebook!;
-		const newStats = await this._resolveStats(this.resource);
-		this._lastResolvedFileStat = newStats;
-		this._register(this._notebook);
-
-		this._name = basename(this._notebook!.uri);
-
-		this._register(this._notebook.onDidChangeContent(() => {
-			this._onDidChangeContent.fire();
-		}));
-		this._register(this._notebook.onDidChangeDirty(() => {
-			this._onDidChangeDirty.fire();
-		}));
-
-		await this._backupFileService.discardBackup(this._workingCopyResource);
-		this._notebook.setDirty(true);
-
-		return this;
-	}
-
 	private async loadFromProvider(forceReloadFromDisk: boolean, editorId: string | undefined, backupId: string | undefined) {
 		const notebook = await this._notebookService.resolveNotebook(this.viewType!, this.resource, forceReloadFromDisk, editorId, backupId);
 		this._notebook = notebook!;

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

@@ -9,7 +9,7 @@ import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/noteb
 import { NotebookExtensionDescription } from 'vs/workbench/api/common/extHost.protocol';
 import { Event } from 'vs/base/common/event';
 import {
-	INotebookTextModel, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, INotebookKernelInfoDto, INotebookTextModelBackup,
+	INotebookTextModel, INotebookRendererInfo, INotebookKernelInfo, INotebookKernelInfoDto,
 	IEditor, ICellEditOperation, NotebookCellOutputsSplice, IOrderedMimeType, IProcessedOutput, INotebookKernelProvider, INotebookKernelInfo2
 } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
@@ -22,12 +22,14 @@ export const INotebookService = createDecorator<INotebookService>('notebookServi
 
 export interface IMainNotebookController {
 	kernel: INotebookKernelInfoDto | undefined;
-	createNotebook(viewType: string, uri: URI, backup: INotebookTextModelBackup | undefined, forceReload: boolean, editorId?: string, backupId?: string): Promise<NotebookTextModel | undefined>;
+	supportBackup: boolean;
+	createNotebook(textModel: NotebookTextModel, editorId?: string, backupId?: string): Promise<void>;
+	reloadNotebook(mainthreadTextModel: NotebookTextModel): Promise<void>;
 	resolveNotebookEditor(viewType: string, uri: URI, editorId: string): Promise<void>;
 	executeNotebookByAttachedKernel(viewType: string, uri: URI, token: CancellationToken): Promise<void>;
 	onDidReceiveMessage(editorId: string, rendererType: string | undefined, message: any): void;
 	executeNotebookCell(uri: URI, handle: number, token: CancellationToken): Promise<void>;
-	removeNotebookDocument(notebook: INotebookTextModel): Promise<void>;
+	removeNotebookDocument(uri: URI): Promise<void>;
 	save(uri: URI, token: CancellationToken): Promise<boolean>;
 	saveAs(uri: URI, target: URI, token: CancellationToken): Promise<boolean>;
 	backup(uri: URI, token: CancellationToken): Promise<string | undefined>;
@@ -58,7 +60,7 @@ export interface INotebookService {
 	getContributedNotebookKernels2(viewType: string, resource: URI, token: CancellationToken): Promise<INotebookKernelInfo2[]>;
 	getRendererInfo(id: string): INotebookRendererInfo | undefined;
 	resolveNotebook(viewType: string, uri: URI, forceReload: boolean, editorId?: string, backupId?: string): Promise<NotebookTextModel | undefined>;
-	createNotebookFromBackup(viewType: string, uri: URI, metadata: NotebookDocumentMetadata, languages: string[], cells: ICellDto2[], editorId?: string): Promise<NotebookTextModel | undefined>;
+	getNotebookTextModel(uri: URI): NotebookTextModel | undefined;
 	executeNotebook(viewType: string, uri: URI, token: CancellationToken): Promise<void>;
 	executeNotebookCell(viewType: string, uri: URI, handle: number, token: CancellationToken): Promise<void>;
 	executeNotebook2(viewType: string, uri: URI, kernelId: string, token: CancellationToken): Promise<void>;

src/vs/workbench/contrib/notebook/test/notebookTextModel.test.ts

@@ -32,7 +32,7 @@ suite('NotebookTextModel', () => {
 				textModel.$applyEdit(textModel.versionId, [
 					{ editType: CellEditType.Insert, index: 1, cells: [new TestCell(viewModel.viewType, 5, ['var e = 5;'], 'javascript', CellKind.Code, [], textModelService)] },
 					{ editType: CellEditType.Insert, index: 3, cells: [new TestCell(viewModel.viewType, 6, ['var f = 6;'], 'javascript', CellKind.Code, [], textModelService)] },
-				], true, true);
+				], true);
 
 				assert.equal(textModel.cells.length, 6);
 
@@ -57,7 +57,7 @@ suite('NotebookTextModel', () => {
 				textModel.$applyEdit(textModel.versionId, [
 					{ editType: CellEditType.Insert, index: 1, cells: [new TestCell(viewModel.viewType, 5, ['var e = 5;'], 'javascript', CellKind.Code, [], textModelService)] },
 					{ editType: CellEditType.Insert, index: 1, cells: [new TestCell(viewModel.viewType, 6, ['var f = 6;'], 'javascript', CellKind.Code, [], textModelService)] },
-				], true, true);
+				], true);
 
 				assert.equal(textModel.cells.length, 6);
 
@@ -82,7 +82,7 @@ suite('NotebookTextModel', () => {
 				textModel.$applyEdit(textModel.versionId, [
 					{ editType: CellEditType.Delete, index: 1, count: 1 },
 					{ editType: CellEditType.Delete, index: 3, count: 1 },
-				], true, true);
+				], true);
 
 				assert.equal(textModel.cells[0].getValue(), 'var a = 1;');
 				assert.equal(textModel.cells[1].getValue(), 'var c = 3;');
@@ -105,7 +105,7 @@ suite('NotebookTextModel', () => {
 				textModel.$applyEdit(textModel.versionId, [
 					{ editType: CellEditType.Delete, index: 1, count: 1 },
 					{ editType: CellEditType.Insert, index: 3, cells: [new TestCell(viewModel.viewType, 5, ['var e = 5;'], 'javascript', CellKind.Code, [], textModelService)] },
-				], true, true);
+				], true);
 
 				assert.equal(textModel.cells.length, 4);
 
@@ -130,7 +130,7 @@ suite('NotebookTextModel', () => {
 				textModel.$applyEdit(textModel.versionId, [
 					{ editType: CellEditType.Delete, index: 1, count: 1 },
 					{ editType: CellEditType.Insert, index: 1, cells: [new TestCell(viewModel.viewType, 5, ['var e = 5;'], 'javascript', CellKind.Code, [], textModelService)] },
-				], true, true);
+				], true);
 
 				assert.equal(textModel.cells.length, 4);
 				assert.equal(textModel.cells[0].getValue(), 'var a = 1;');