diff --git a/src/vs/workbench/contrib/notebook/browser/contrib/undoRedo/notebookUndoRedo.ts b/src/vs/workbench/contrib/notebook/browser/contrib/undoRedo/notebookUndoRedo.ts
index be1a847ec81..93d8764b970 100644
--- a/src/vs/workbench/contrib/notebook/browser/contrib/undoRedo/notebookUndoRedo.ts
+++ b/src/vs/workbench/contrib/notebook/browser/contrib/undoRedo/notebookUndoRedo.ts
@@ -11,6 +11,7 @@ import { CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { CellEditState, getNotebookEditorFromEditorPane } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { RedoCommand, UndoCommand } from 'vs/editor/browser/editorExtensions';
+import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
 
 class NotebookUndoRedoContribution extends Disposable {
 
@@ -26,38 +27,46 @@ class NotebookUndoRedoContribution extends Disposable {
 		 * We now make the notebook undo/redo impl the highest priority so we don't skip the "invalid" undo/redo element in the same notebook document
 		 */
 		this._register(UndoCommand.addImplementation(10000 + 5, 'notebook-undo-redo', () => {
-			const editor = getNotebookEditorFromEditorPane(this._editorService.activeEditorPane);
-			if (editor?.viewModel) {
-				return editor.viewModel.undo().then(cellResources => {
-					if (cellResources?.length) {
-						editor?.viewModel?.viewCells.forEach(cell => {
-							if (cell.cellKind === CellKind.Markup && cellResources.find(resource => resource.fragment === cell.model.uri.fragment)) {
-								cell.updateEditState(CellEditState.Editing, 'undo');
-							}
-						});
+			const editor = getNotebookEditorFromEditorPane(this._editorService.activeEditorPane) as NotebookEditorWidget | undefined;
+			if (editor?.hasModel()) {
+				const activeCodeEditor = editor.activeCodeEditor;
 
-						editor?.setOptions({ cellOptions: { resource: cellResources[0] }, preserveFocus: true });
-					}
-				});
+				if (activeCodeEditor && activeCodeEditor.hasTextFocus()) {
+					return editor.viewModel.undo().then(cellResources => {
+						if (cellResources?.length) {
+							editor?.viewModel?.viewCells.forEach(cell => {
+								if (cell.cellKind === CellKind.Markup && cellResources.find(resource => resource.fragment === cell.model.uri.fragment)) {
+									cell.updateEditState(CellEditState.Editing, 'undo');
+								}
+							});
+
+							editor?.setOptions({ cellOptions: { resource: cellResources[0] }, preserveFocus: true });
+						}
+					});
+				}
 			}
 
 			return false;
 		}));
 
 		this._register(RedoCommand.addImplementation(10000 + 5, 'notebook-undo-redo', () => {
-			const editor = getNotebookEditorFromEditorPane(this._editorService.activeEditorPane);
-			if (editor?.viewModel) {
-				return editor.viewModel.redo().then(cellResources => {
-					if (cellResources?.length) {
-						editor?.viewModel?.viewCells.forEach(cell => {
-							if (cell.cellKind === CellKind.Markup && cellResources.find(resource => resource.fragment === cell.model.uri.fragment)) {
-								cell.updateEditState(CellEditState.Editing, 'redo');
-							}
-						});
+			const editor = getNotebookEditorFromEditorPane(this._editorService.activeEditorPane) as NotebookEditorWidget | undefined;
+			if (editor?.hasModel()) {
+				const activeCodeEditor = editor.activeCodeEditor;
 
-						editor?.setOptions({ cellOptions: { resource: cellResources[0] }, preserveFocus: true });
-					}
-				});
+				if (activeCodeEditor && activeCodeEditor.hasTextFocus()) {
+					return editor.viewModel.redo().then(cellResources => {
+						if (cellResources?.length) {
+							editor?.viewModel?.viewCells.forEach(cell => {
+								if (cell.cellKind === CellKind.Markup && cellResources.find(resource => resource.fragment === cell.model.uri.fragment)) {
+									cell.updateEditState(CellEditState.Editing, 'redo');
+								}
+							});
+
+							editor?.setOptions({ cellOptions: { resource: cellResources[0] }, preserveFocus: true });
+						}
+					});
+				}
 			}
 
 			return false;