microsoft · rchiodo · Feb 8, 2020 · Feb 7, 2020 · Feb 8, 2020 · Feb 8, 2020
diff --git a/src/client/datascience/interactive-ipynb/nativeEditor.ts b/src/client/datascience/interactive-ipynb/nativeEditor.ts
@@ -45,6 +45,9 @@ import {
     WebViewViewChangeEventArgs
 } from '../types';
 
+// tslint:disable-next-line: no-require-imports
+import cloneDeep = require('lodash/cloneDeep');
+
 const nativeEditorDir = path.join(EXTENSION_ROOT_DIR, 'out', 'datascience-ui', 'notebook');
 @injectable()
 export class NativeEditor extends InteractiveBase implements INotebookEditor {
@@ -345,7 +348,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
                 source: 'user',
                 kind: 'modify',
                 newCells: modified,
-                oldCells: unmodified,
+                oldCells: cloneDeep(unmodified),
                 oldDirty: this._model.isDirty,
                 newDirty: true
             });

diff --git a/src/client/datascience/interactive-ipynb/nativeEditorProvider.ts b/src/client/datascience/interactive-ipynb/nativeEditorProvider.ts
@@ -197,6 +197,12 @@ export class NativeEditorProvider implements INotebookEditorProvider, WebviewCus
 
     private closedEditor(editor: INotebookEditor): void {
         this.openedEditors.delete(editor);
+        // If last editor, dispose of the storage
+        const key = editor.file.toString();
+        if (![...this.openedEditors].find(e => e.file.toString() === key)) {
+            this.modelChangedHandlers.delete(key);
+            this.models.delete(key);
+        }
         this._onDidCloseNotebookEditor.fire(editor);
     }
 

diff --git a/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts b/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts
@@ -26,14 +26,15 @@ const NotebookTransferKey = 'notebook-transfered';
 interface INativeEditorStorageState {
     file: Uri;
     cells: ICell[];
-    changeCountSinceSave: number;
+    changeCount: number;
+    saveChangeCount: number;
     notebookJson: Partial<nbformat.INotebookContent>;
 }
 
 @injectable()
 export class NativeEditorStorage implements INotebookModel, INotebookStorage {
     public get isDirty(): boolean {
-        return this._state.changeCountSinceSave > 0;
+        return this._state.changeCount !== this._state.saveChangeCount && this._state.changeCount !== 0;
     }
     public get changed(): Event<NotebookModelChange> {
         return this._changedEmitter.event;
@@ -49,7 +50,7 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
         return this._state.cells;
     }
     private _changedEmitter = new EventEmitter<NotebookModelChange>();
-    private _state: INativeEditorStorageState = { file: Uri.file(''), changeCountSinceSave: 0, cells: [], notebookJson: {} };
+    private _state: INativeEditorStorageState = { file: Uri.file(''), changeCount: 0, saveChangeCount: 0, cells: [], notebookJson: {} };
     private indentAmount: string = ' ';
 
     constructor(
@@ -145,7 +146,7 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
             case 'file':
                 changed = !this.fileSystem.arePathsSame(this._state.file.fsPath, change.newFile.fsPath);
                 this._state.file = change.newFile;
-                this._state.changeCountSinceSave = 0;
+                this._state.saveChangeCount = this._state.changeCount;
                 break;
             default:
                 break;
@@ -154,7 +155,7 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
         // Dirty state comes from undo. At least VS code will track it that way. However
         // skip version and file changes as we don't forward those to VS code
         if (change.kind !== 'file' && change.kind !== 'version') {
-            this._state.changeCountSinceSave += 1;
+            this._state.changeCount += 1;
         }
 
         return changed;
@@ -194,7 +195,7 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
         // Dirty state comes from undo. At least VS code will track it that way.
         // Note unlike redo, 'file' and 'version' are not possible on undo as
         // we don't send them to VS code.
-        this._state.changeCountSinceSave -= 1;
+        this._state.changeCount -= 1;
 
         return changed;
     }

diff --git a/src/datascience-ui/native-editor/redux/reducers/creation.ts b/src/datascience-ui/native-editor/redux/reducers/creation.ts
@@ -278,8 +278,16 @@ export namespace Creation {
                     ...disabledQueueArg,
                     payload: { ...arg.payload, data: { firstCellId: arg.payload.data.secondCellId, secondCellId: arg.payload.data.firstCellId } }
                 });
+            case 'modify':
+                // Undo for modify should reapply the outputs. Go through each and apply the update
+                let result = arg.prevState;
+                arg.payload.data.oldCells.forEach(c => {
+                    result = updateCell({ ...disabledQueueArg, prevState: result, payload: { ...arg.payload, data: c } });
+                });
+                return result;
+
             default:
-                // Modify, file, version can all be ignored.
+                // File, version can be ignored.
                 break;
         }
 
@@ -307,6 +315,13 @@ export namespace Creation {
                     ...disabledQueueArg,
                     payload: { ...arg.payload, data: { firstCellId: arg.payload.data.secondCellId, secondCellId: arg.payload.data.firstCellId } }
                 });
+            case 'modify':
+                // Redo for modify should reapply the outputs. Go through each and apply the update
+                let result = arg.prevState;
+                arg.payload.data.newCells.forEach(c => {
+                    result = updateCell({ ...disabledQueueArg, prevState: result, payload: { ...arg.payload, data: c } });
+                });
+                return result;
             default:
                 // Modify, file, version can all be ignored.
                 break;