diff --git a/src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts b/src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts
index 136cfd3b622b..7a9953b2be87 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts
@@ -30,7 +30,7 @@ import { NotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookEd
 import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/browser/notebookEditorInput';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { NotebookService } from 'vs/workbench/contrib/notebook/browser/notebookServiceImpl';
-import { CellKind, CellUri, NotebookDocumentBackupData, NotebookEditorPriority } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellKind, CellUri, getCellUndoRedoComparisonKey, NotebookDocumentBackupData, NotebookEditorPriority } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
 import { IEditorGroup } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { IEditorService, IOpenEditorOverride } from 'vs/workbench/services/editor/common/editorService';
@@ -151,12 +151,7 @@ export class NotebookContribution extends Disposable implements IWorkbenchContri
 
 		this._register(undoRedoService.registerUriComparisonKeyComputer(CellUri.scheme, {
 			getComparisonKey: (uri: URI): string => {
-				const data = CellUri.parse(uri);
-				if (!data) {
-					return uri.toString();
-				}
-
-				return data.notebook.toString();
+				return getCellUndoRedoComparisonKey(uri);
 			}
 		}));
 
diff --git a/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts b/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts
index 057122d125b9..1e5ab3750b3a 100644
--- a/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts
+++ b/src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts
@@ -32,6 +32,7 @@ import { IPosition, Position } from 'vs/editor/common/core/position';
 import { SplitCellEdit, JoinCellEdit } from 'vs/workbench/contrib/notebook/browser/viewModel/cellEdit';
 import { BaseCellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/baseCellViewModel';
 import { PieceTreeTextBuffer } from 'vs/editor/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer';
+import { MultiModelEditStackElement, SingleModelEditStackElement } from 'vs/editor/common/model/editStack';
 
 export interface INotebookEditorViewState {
 	editingCells: { [key: number]: boolean };
@@ -1051,12 +1052,38 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		});
 	}
 
+	async withElement(element: SingleModelEditStackElement | MultiModelEditStackElement, callback: () => Promise<void>) {
+		const viewCells = this._viewCells.filter(cell => element.matchesResource(cell.uri));
+		const refs = await Promise.all(viewCells.map(cell => cell.model.resolveTextModelRef()));
+		await callback();
+		refs.forEach(ref => ref.dispose());
+	}
+
 	async undo() {
+		const editStack = this._undoService.getElements(this.uri);
+		const element = editStack.past.length ? editStack.past[editStack.past.length - 1] : undefined;
+
+		if (element && element instanceof SingleModelEditStackElement || element instanceof MultiModelEditStackElement) {
+			return await this.withElement(element, async () => {
+				await this._undoService.undo(this.uri);
+			});
+		}
+
 		await this._undoService.undo(this.uri);
 	}
 
 	async redo() {
+		const editStack = this._undoService.getElements(this.uri);
+		const element = editStack.future[0];
+
+		if (element && element instanceof SingleModelEditStackElement || element instanceof MultiModelEditStackElement) {
+			return await this.withElement(element, async () => {
+				await this._undoService.redo(this.uri);
+			});
+		}
+
 		await this._undoService.redo(this.uri);
+
 	}
 
 	equal(notebook: NotebookTextModel) {
diff --git a/src/vs/workbench/contrib/notebook/common/notebookCommon.ts b/src/vs/workbench/contrib/notebook/common/notebookCommon.ts
index 99d94fc20f4c..2bddcefc19c5 100644
--- a/src/vs/workbench/contrib/notebook/common/notebookCommon.ts
+++ b/src/vs/workbench/contrib/notebook/common/notebookCommon.ts
@@ -413,6 +413,15 @@ export interface NotebookDataDto {
 	readonly metadata: NotebookDocumentMetadata;
 }
 
+export function getCellUndoRedoComparisonKey(uri: URI) {
+	const data = CellUri.parse(uri);
+	if (!data) {
+		return uri.toString();
+	}
+
+	return data.notebook.toString();
+}
+
 
 export namespace CellUri {