Fix #97694.

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

@@ -8,7 +8,7 @@ import { MainContext, MainThreadNotebookShape, NotebookExtensionDescription, IEx
 import { Disposable } from 'vs/base/common/lifecycle';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { INotebookService, IMainNotebookController } from 'vs/workbench/contrib/notebook/common/notebookService';
-import { INotebookTextModel, INotebookMimeTypeSelector, NOTEBOOK_DISPLAY_ORDER, NotebookCellOutputsSplice, NotebookDocumentMetadata, NotebookCellMetadata, ICellEditOperation, ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellEditType, CellKind, INotebookKernelInfo, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookMimeTypeSelector, NOTEBOOK_DISPLAY_ORDER, NotebookCellOutputsSplice, NotebookDocumentMetadata, NotebookCellMetadata, ICellEditOperation, ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellEditType, CellKind, INotebookKernelInfo, INotebookKernelInfoDto, INotebookTextModelBackup } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
@@ -210,7 +210,7 @@ export class MainThreadNotebookController implements IMainNotebookController {
 	) {
 	}
 
-	async createNotebook(viewType: string, uri: URI, forBackup: boolean, forceReload: boolean): Promise<NotebookTextModel | undefined> {
+	async createNotebook(viewType: string, uri: URI, backup: INotebookTextModelBackup | undefined, forceReload: boolean): Promise<NotebookTextModel | undefined> {
 		let mainthreadNotebook = this._mapping.get(URI.from(uri).toString());
 
 		if (mainthreadNotebook) {
@@ -231,9 +231,41 @@ export class MainThreadNotebookController implements IMainNotebookController {
 		}
 
 		let document = new MainThreadNotebookDocument(this._proxy, MainThreadNotebookController.documentHandle++, viewType, uri, generateUuid());
-		await this.createNotebookDocument(document);
+		this._mapping.set(document.uri.toString(), document);
+
+		if (backup) {
+			// trigger events
+			document.textModel.metadata = backup.metadata;
+			document.textModel.languages = backup.languages;
+
+			document.textModel.applyEdit(document.textModel.versionId, [
+				{
+					editType: CellEditType.Insert,
+					index: 0,
+					cells: backup.cells
+				}
+			]);
+
+			await this._proxy.$acceptDocumentAndEditorsDelta({
+				addedDocuments: [{
+					viewType: document.viewType,
+					handle: document.handle,
+					webviewId: document.webviewId,
+					uri: document.uri,
+					metadata: document.textModel.metadata,
+					versionId: document.textModel.versionId,
+					cells: document.textModel.cells.map(cell => ({
+						handle: cell.handle,
+						uri: cell.uri,
+						source: cell.textBuffer.getLinesContent(),
+						language: cell.language,
+						cellKind: cell.cellKind,
+						outputs: cell.outputs,
+						metadata: cell.metadata
+					}))
+				}]
+			});
 
-		if (forBackup) {
 			return document.textModel;
 		}
 
@@ -253,6 +285,26 @@ export class MainThreadNotebookController implements IMainNotebookController {
 			document.textModel.insertTemplateCell(mainCell);
 		}
 
+		await this._proxy.$acceptDocumentAndEditorsDelta({
+			addedDocuments: [{
+				viewType: document.viewType,
+				handle: document.handle,
+				webviewId: document.webviewId,
+				uri: document.uri,
+				metadata: document.textModel.metadata,
+				versionId: document.textModel.versionId,
+				cells: document.textModel.cells.map(cell => ({
+					handle: cell.handle,
+					uri: cell.uri,
+					source: cell.textBuffer.getLinesContent(),
+					language: cell.language,
+					cellKind: cell.cellKind,
+					outputs: cell.outputs,
+					metadata: cell.metadata
+				}))
+			}]
+		});
+
 		this._proxy.$acceptEditorPropertiesChanged(uri, { selections: null, metadata: document.textModel.metadata });
 
 		return document.textModel;
@@ -283,20 +335,6 @@ export class MainThreadNotebookController implements IMainNotebookController {
 		this._proxy.$onDidReceiveMessage(uri, message);
 	}
 
-	async createNotebookDocument(document: MainThreadNotebookDocument): Promise<void> {
-		this._mapping.set(document.uri.toString(), document);
-
-		await this._proxy.$acceptDocumentAndEditorsDelta({
-			addedDocuments: [{
-				viewType: document.viewType,
-				handle: document.handle,
-				webviewId: document.webviewId,
-				uri: document.uri,
-				metadata: document.textModel.metadata
-			}]
-		});
-	}
-
 	async removeNotebookDocument(notebook: INotebookTextModel): Promise<void> {
 		let document = this._mapping.get(URI.from(notebook.uri).toString());
 

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

@@ -51,7 +51,7 @@ import { TunnelDto } from 'vs/workbench/api/common/extHostTunnelService';
 import { TunnelOptions } from 'vs/platform/remote/common/tunnel';
 import { Timeline, TimelineChangeEvent, TimelineOptions, TimelineProviderDescriptor, InternalTimelineOptions } from 'vs/workbench/contrib/timeline/common/timeline';
 import { revive } from 'vs/base/common/marshalling';
-import { INotebookMimeTypeSelector, IOutput, INotebookDisplayOrder, NotebookCellMetadata, NotebookDocumentMetadata, ICellEditOperation, NotebookCellsChangedEvent, NotebookDataDto, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookMimeTypeSelector, IOutput, INotebookDisplayOrder, NotebookCellMetadata, NotebookDocumentMetadata, ICellEditOperation, NotebookCellsChangedEvent, NotebookDataDto, INotebookKernelInfoDto, IMainCellDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { CallHierarchyItem } from 'vs/workbench/contrib/callHierarchy/common/callHierarchy';
 import { Dto } from 'vs/base/common/types';
 import { ISerializableEnvironmentVariableCollection } from 'vs/workbench/contrib/terminal/common/environmentVariable';
@@ -1547,7 +1547,8 @@ export interface INotebookModelAddedData {
 	uri: UriComponents;
 	handle: number;
 	webviewId: string;
-	// versionId: number;
+	versionId: number;
+	cells: IMainCellDto[],
 	viewType: string;
 	metadata?: NotebookDocumentMetadata;
 }

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

@@ -633,6 +633,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 	private readonly _notebookContentProviders = new Map<string, { readonly provider: vscode.NotebookContentProvider, readonly extension: IExtensionDescription; }>();
 	private readonly _notebookKernels = new Map<string, { readonly kernel: vscode.NotebookKernel, readonly extension: IExtensionDescription; }>();
 	private readonly _documents = new Map<string, ExtHostNotebookDocument>();
+	private readonly _unInitializedDocuments = new Map<string, ExtHostNotebookDocument>();
 	private readonly _editors = new Map<string, { editor: ExtHostNotebookEditor, onDidReceiveMessage: Emitter<any>; }>();
 	private readonly _notebookOutputRenderers = new Map<number, ExtHostNotebookOutputRenderer>();
 
@@ -745,10 +746,17 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 	}
 
 	async $resolveNotebookData(viewType: string, uri: UriComponents): Promise<NotebookDataDto | undefined> {
-		let provider = this._notebookContentProviders.get(viewType);
-		let document = this._documents.get(URI.revive(uri).toString());
+		const provider = this._notebookContentProviders.get(viewType);
+		const revivedUri = URI.revive(uri);
+
+		if (provider) {
+			let document = this._documents.get(URI.revive(uri).toString());
+
+			if (!document) {
+				document = this._unInitializedDocuments.get(revivedUri.toString()) ?? new ExtHostNotebookDocument(this._proxy, this._documentsAndEditors, viewType, revivedUri, this);
+				this._unInitializedDocuments.set(revivedUri.toString(), document);
+			}
 
-		if (provider && document) {
 			const rawCells = await provider.provider.openNotebook(URI.revive(uri));
 			const renderers = new Set<number>();
 			const dto = {
@@ -997,9 +1005,11 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		if (delta.addedDocuments) {
 			delta.addedDocuments.forEach(modelData => {
 				const revivedUri = URI.revive(modelData.uri);
+				const revivedUriStr = revivedUri.toString();
 				const viewType = modelData.viewType;
-				if (!this._documents.has(revivedUri.toString())) {
-					let document = new ExtHostNotebookDocument(this._proxy, this._documentsAndEditors, viewType, revivedUri, this);
+				if (!this._documents.has(revivedUriStr)) {
+					let document = this._unInitializedDocuments.get(revivedUriStr) ?? new ExtHostNotebookDocument(this._proxy, this._documentsAndEditors, viewType, revivedUri, this);
+					this._unInitializedDocuments.delete(revivedUriStr);
 					if (modelData.metadata) {
 						document.metadata = {
 							...notebookDocumentMetadataDefaults,
@@ -1007,11 +1017,23 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 						};
 					}
 
-					this._documents.set(revivedUri.toString(), document);
+					document.accpetModelChanged({
+						kind: NotebookCellsChangeType.ModelChange,
+						versionId: modelData.versionId,
+						changes: [
+							[
+								0,
+								0,
+								modelData.cells
+							]
+						]
+					});
+
+					this._documents.set(revivedUriStr, document);
 				}
 
 				const onDidReceiveMessage = new Emitter<any>();
-				const document = this._documents.get(revivedUri.toString())!;
+				const document = this._documents.get(revivedUriStr)!;
 
 				let editor = new ExtHostNotebookEditor(
 					viewType,
@@ -1028,7 +1050,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 				this._onDidOpenNotebookDocument.fire(document);
 
 				// TODO, does it already exist?
-				this._editors.set(revivedUri.toString(), { editor, onDidReceiveMessage });
+				this._editors.set(revivedUriStr, { editor, onDidReceiveMessage });
 			});
 		}
 

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

@@ -11,7 +11,7 @@ import { notebookProviderExtensionPoint, notebookRendererExtensionPoint } from '
 import { NotebookProviderInfo } 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, INotebookMimeTypeSelector, INotebookRendererInfo, NotebookDocumentMetadata, CellEditType, ICellDto2, INotebookKernelInfo } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookMimeTypeSelector, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo } 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';
@@ -264,7 +264,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 			return undefined;
 		}
 
-		const notebookModel = await provider.controller.createNotebook(viewType, uri, true, false);
+		const notebookModel = await provider.controller.createNotebook(viewType, uri, { metadata, languages, cells }, false);
 		if (!notebookModel) {
 			return undefined;
 		}
@@ -276,18 +276,6 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 			(model) => this._onWillDispose(model),
 		);
 		this._models[modelId] = modelData;
-
-		notebookModel.metadata = metadata;
-		notebookModel.languages = languages;
-
-		notebookModel.applyEdit(notebookModel.versionId, [
-			{
-				editType: CellEditType.Insert,
-				index: 0,
-				cells: cells
-			}
-		]);
-
 		return modelData.model;
 	}
 
@@ -299,7 +287,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 
 		let notebookModel: NotebookTextModel | undefined;
 
-		notebookModel = await provider.controller.createNotebook(viewType, uri, false, forceReload);
+		notebookModel = await provider.controller.createNotebook(viewType, uri, undefined, forceReload);
 
 		// new notebook model created
 		const modelId = MODEL_ID(uri);

src/vs/workbench/contrib/notebook/browser/view/notebookCellList.ts

@@ -357,7 +357,7 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 			}
 		});
 
-		if (!selectionsLeft.length && this._viewModel!.viewCells) {
+		if (!selectionsLeft.length && this._viewModel!.viewCells.length) {
 			// after splice, the selected cells are deleted
 			this._viewModel!.selectionHandles = [this._viewModel!.viewCells[0].handle];
 		}

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

@@ -502,3 +502,8 @@ export interface INotebookEditorModel extends IEditorModel {
 	save(): Promise<boolean>;
 }
 
+export interface INotebookTextModelBackup {
+	metadata: NotebookDocumentMetadata;
+	languages: string[];
+	cells: ICellDto2[]
+}

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

@@ -8,7 +8,7 @@ import { URI } from 'vs/base/common/uri';
 import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
 import { NotebookExtensionDescription } from 'vs/workbench/api/common/extHost.protocol';
 import { Event } from 'vs/base/common/event';
-import { INotebookTextModel, INotebookMimeTypeSelector, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookMimeTypeSelector, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, INotebookKernelInfoDto, INotebookTextModelBackup } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
@@ -18,7 +18,7 @@ export const INotebookService = createDecorator<INotebookService>('notebookServi
 
 export interface IMainNotebookController {
 	kernel: INotebookKernelInfoDto | undefined;
-	createNotebook(viewType: string, uri: URI, forBackup: boolean, forceReload: boolean): Promise<NotebookTextModel | undefined>;
+	createNotebook(viewType: string, uri: URI, backup: INotebookTextModelBackup | undefined, forceReload: boolean): Promise<NotebookTextModel | undefined>;
 	executeNotebook(viewType: string, uri: URI, useAttachedKernel: boolean, token: CancellationToken): Promise<void>;
 	onDidReceiveMessage(uri: URI, message: any): void;
 	executeNotebookCell(uri: URI, handle: number, useAttachedKernel: boolean, token: CancellationToken): Promise<void>;

src/vs/workbench/test/browser/api/extHostNotebookConcatDocument.test.ts

@@ -53,21 +53,18 @@ suite('NotebookConcatDocument', function () {
 				handle: 0,
 				uri: notebookUri,
 				viewType: 'test',
-				webviewId: 'testid'
+				webviewId: 'testid',
+				cells: [{
+					handle: 0,
+					uri: CellUri.generate(notebookUri, 0),
+					source: ['### Heading'],
+					language: 'markdown',
+					cellKind: CellKind.Markdown,
+					outputs: [],
+				}],
+				versionId: 0
 			}]
 		});
-		extHostNotebooks.$acceptModelChanged(notebookUri, {
-			kind: NotebookCellsChangeType.ModelChange,
-			versionId: 0,
-			changes: [[0, 0, [{
-				handle: 0,
-				uri: CellUri.generate(notebookUri, 0),
-				source: ['### Heading'],
-				language: 'markdown',
-				cellKind: CellKind.Markdown,
-				outputs: [],
-			}]]]
-		});
 		await extHostNotebooks.$acceptDocumentAndEditorsDelta({ newActiveEditor: notebookUri });
 
 		notebook = extHostNotebooks.activeNotebookDocument!;