From da88f129d3c555ccfe2de7629f9bd82eb7985365 Mon Sep 17 00:00:00 2001
From: rebornix <penn.lv@gmail.com>
Date: Wed, 29 Apr 2020 10:39:27 -0700
Subject: [PATCH] Re #96398. Fix the data loss when switching

---
 src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts  | 1 +
 src/vs/workbench/contrib/notebook/browser/notebookEditor.ts   | 2 ++
 .../workbench/contrib/notebook/browser/notebookEditorInput.ts | 4 ++++
 .../contrib/notebook/browser/viewModel/baseCellViewModel.ts   | 2 ++
 4 files changed, 9 insertions(+)

diff --git a/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts b/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts
index 972f86ec457..f37c2ef728c 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts
@@ -102,6 +102,7 @@ export interface ICellViewModel {
 	currentTokenSource: CancellationTokenSource | undefined;
 	focusMode: CellFocusMode;
 	getText(): string;
+	save(): void;
 	metadata: NotebookCellMetadata | undefined;
 	getEvaluatedMetadata(documentMetadata: NotebookDocumentMetadata | undefined): NotebookCellMetadata;
 }
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditor.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditor.ts
index 07d97f4bfa8..17a585c3774 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditor.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditor.ts
@@ -636,6 +636,8 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 
 			state.contributionsState = contributionsState;
 			this.editorMemento.saveEditorState(this.group, input.resource, state);
+
+			this.notebookViewModel.viewCells.forEach(cell => cell.save());
 		}
 	}
 
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorInput.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorInput.ts
index 7f53d24bdda..600c5aaebf2 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditorInput.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditorInput.ts
@@ -173,6 +173,10 @@ export class NotebookEditorInput extends EditorInput {
 	}
 
 	async resolve(): Promise<NotebookEditorModel> {
+		if (this.textModel) {
+			return this.textModel;
+		}
+
 		if (!this.promise) {
 			if (!await this.notebookService.canResolve(this.viewType!)) {
 				throw new Error(`Cannot open notebook of type '${this.viewType}'`);
diff --git a/src/vs/workbench/contrib/notebook/browser/viewModel/baseCellViewModel.ts b/src/vs/workbench/contrib/notebook/browser/viewModel/baseCellViewModel.ts
index 6779eb7c8d4..63fa4781374 100644
--- a/src/vs/workbench/contrib/notebook/browser/viewModel/baseCellViewModel.ts
+++ b/src/vs/workbench/contrib/notebook/browser/viewModel/baseCellViewModel.ts
@@ -199,6 +199,8 @@ export abstract class BaseCellViewModel extends Disposable implements ICellViewM
 		return this.model.source.join('\n');
 	}
 
+	abstract save(): void;
+
 	private saveViewState(): editorCommon.ICodeEditorViewState | null {
 		if (!this._textEditor) {
 			return null;