hook native undo/redo with notebook

extensions/vscode-notebook-tests/src/notebook.test.ts

@@ -663,7 +663,7 @@ suite('notebook undo redo', () => {
 
 
 		// undo should bring back the deleted cell, and revert to previous content and selection
-		await vscode.commands.executeCommand('notebook.undo');
+		await vscode.commands.executeCommand('undo');
 		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.length, 3);
 		assert.equal(vscode.notebook.activeNotebookEditor!.document.cells.indexOf(vscode.notebook.activeNotebookEditor!.selection!), 1);
 		assert.equal(vscode.notebook.activeNotebookEditor?.selection?.document.getText(), 'var abc = 0;');
@@ -729,7 +729,7 @@ suite('notebook undo redo', () => {
 		assert.equal(cellOutputsAddedRet.cells[0].outputs.length, 1);
 
 		const cellOutputClear = getEventOncePromise<vscode.NotebookCellOutputsChangeEvent>(vscode.notebook.onDidChangeCellOutputs);
-		await vscode.commands.executeCommand('notebook.undo');
+		await vscode.commands.executeCommand('undo');
 		const cellOutputsCleardRet = await cellOutputClear;
 		assert.deepEqual(cellOutputsCleardRet, {
 			document: vscode.notebook.activeNotebookEditor!.document,

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

@@ -30,8 +30,6 @@ const EXECUTE_NOTEBOOK_COMMAND_ID = 'notebook.execute';
 const CANCEL_NOTEBOOK_COMMAND_ID = 'notebook.cancelExecution';
 const NOTEBOOK_FOCUS_TOP = 'notebook.focusTop';
 const NOTEBOOK_FOCUS_BOTTOM = 'notebook.focusBottom';
-const NOTEBOOK_REDO = 'notebook.redo';
-const NOTEBOOK_UNDO = 'notebook.undo';
 const NOTEBOOK_FOCUS_PREVIOUS_EDITOR = 'notebook.focusPreviousEditor';
 const NOTEBOOK_FOCUS_NEXT_EDITOR = 'notebook.focusNextEditor';
 const CLEAR_ALL_CELLS_OUTPUTS_COMMAND_ID = 'notebook.clearAllCellsOutputs';
@@ -1108,42 +1106,6 @@ registerAction2(class extends NotebookCellAction {
 });
 
 
-registerAction2(class extends NotebookAction {
-	constructor() {
-		super({
-			id: NOTEBOOK_UNDO,
-			title: localize('undo', 'Undo'),
-			keybinding: {
-				when: ContextKeyExpr.and(NOTEBOOK_EDITOR_FOCUSED, ContextKeyExpr.not(InputFocusedContextKey)),
-				primary: KeyMod.CtrlCmd | KeyCode.KEY_Z,
-				weight: KeybindingWeight.WorkbenchContrib
-			}
-		});
-	}
-
-	async runWithContext(accessor: ServicesAccessor, context: INotebookActionContext): Promise<void> {
-		await context.notebookEditor.viewModel?.undo();
-	}
-});
-
-registerAction2(class extends NotebookAction {
-	constructor() {
-		super({
-			id: NOTEBOOK_REDO,
-			title: localize('redo', 'Redo'),
-			keybinding: {
-				when: ContextKeyExpr.and(NOTEBOOK_EDITOR_FOCUSED, ContextKeyExpr.not(InputFocusedContextKey)),
-				primary: KeyMod.CtrlCmd | KeyMod.Shift | KeyCode.KEY_Z,
-				weight: KeybindingWeight.WorkbenchContrib
-			}
-		});
-	}
-
-	async runWithContext(accessor: ServicesAccessor, context: INotebookActionContext): Promise<void> {
-		await context.notebookEditor.viewModel?.redo();
-	}
-});
-
 registerAction2(class extends NotebookAction {
 	constructor() {
 		super({

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

@@ -26,6 +26,7 @@ import { CellKind, IProcessedOutput, IRenderOutput, NotebookCellMetadata, Notebo
 import { Webview } from 'vs/workbench/contrib/webview/browser/webview';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { IMenu } from 'vs/platform/actions/common/actions';
+import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 
 export const KEYBINDING_CONTEXT_NOTEBOOK_FIND_WIDGET_FOCUSED = new RawContextKey<boolean>('notebookFindWidgetFocused', false);
 
@@ -643,3 +644,9 @@ export function getVisibleCells(cells: CellViewModel[], hiddenRanges: ICellRange
 
 	return result;
 }
+
+export function getActiveNotebookEditor(editorService: IEditorService): INotebookEditor | undefined {
+	// TODO can `isNotebookEditor` be on INotebookEditor to avoid a circular dependency?
+	const activeEditorPane = editorService.activeEditorPane as unknown as { isNotebookEditor?: boolean } | undefined;
+	return activeEditorPane?.isNotebookEditor ? (editorService.activeEditorPane?.getControl() as INotebookEditor) : undefined;
+}

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

@@ -21,7 +21,7 @@ import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/mode
 import { INotebookService, IMainNotebookController } from 'vs/workbench/contrib/notebook/common/notebookService';
 import * as glob from 'vs/base/common/glob';
 import { basename } from 'vs/base/common/path';
-import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+import { getActiveNotebookEditor, INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IAccessibilityService } from 'vs/platform/accessibility/common/accessibility';
 import { Memento } from 'vs/workbench/common/memento';
@@ -32,6 +32,7 @@ import { flatten } from 'vs/base/common/arrays';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { NotebookKernelProviderAssociationRegistry, updateNotebookKernelProvideAssociationSchema, NotebookViewTypesExtensionRegistry } from 'vs/workbench/contrib/notebook/browser/notebookKernelAssociation';
 import { PureNotebookOutputRenderer } from 'vs/workbench/contrib/notebook/browser/notebookPureOutputRenderer';
+import { RedoCommand, UndoCommand } from 'vs/editor/browser/editorExtensions';
 
 function MODEL_ID(resource: URI): string {
 	return resource.toString();
@@ -321,6 +322,27 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		this._register(this._accessibilityService.onDidChangeScreenReaderOptimized(() => {
 			updateOrder();
 		}));
+
+		const PRIORITY = 105;
+		this._register(UndoCommand.addImplementation(PRIORITY, () => {
+			const editor = getActiveNotebookEditor(this._editorService);
+			if (editor?.viewModel) {
+				editor?.viewModel.undo();
+				return true;
+			}
+
+			return false;
+		}));
+
+		this._register(RedoCommand.addImplementation(PRIORITY, () => {
+			const editor = getActiveNotebookEditor(this._editorService);
+			if (editor?.viewModel) {
+				editor?.viewModel.redo();
+				return true;
+			}
+
+			return false;
+		}));
 	}
 
 	getViewTypes(): ICustomEditorInfo[] {