Improve Undo and Redo functionality for notebook Rich Text editors (#15627)

src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts

@@ -74,9 +74,9 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 			if ((e.ctrlKey || e.metaKey) && e.key === 'a') {
 				preventDefaultAndExecCommand(e, 'selectAll');
 			} else if ((e.metaKey && e.shiftKey && e.key === 'z') || (e.ctrlKey && e.key === 'y') && !this.markdownMode) {
-				preventDefaultAndExecCommand(e, 'redo');
+				this.redoRichTextChange();
 			} else if ((e.ctrlKey || e.metaKey) && e.key === 'z') {
-				preventDefaultAndExecCommand(e, 'undo');
+				this.undoRichTextChange();
 			} else if (e.shiftKey && e.key === 'Tab') {
 				preventDefaultAndExecCommand(e, 'outdent');
 			} else if (e.key === 'Tab') {
@@ -104,12 +104,15 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 	private _highlightRange: NotebookRange;
 	private _isFindActive: boolean = false;
 
+	private readonly _undoStack = new RichTextEditStack();
+	private readonly _redoStack = new RichTextEditStack();
+
 	constructor(
 		@Inject(forwardRef(() => ChangeDetectorRef)) private _changeRef: ChangeDetectorRef,
 		@Inject(IInstantiationService) private _instantiationService: IInstantiationService,
 		@Inject(IWorkbenchThemeService) private themeService: IWorkbenchThemeService,
 		@Inject(IConfigurationService) private _configurationService: IConfigurationService,
-		@Inject(INotebookService) private _notebookService: INotebookService,
+		@Inject(INotebookService) private _notebookService: INotebookService
 	) {
 		super();
 		this.markdownRenderer = this._instantiationService.createInstance(NotebookMarkdownRenderer);
@@ -245,6 +248,8 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 			if (this._previewMode) {
 				let outputElement = <HTMLElement>this.output.nativeElement;
 				outputElement.innerHTML = this.markdownResult.element.innerHTML;
+				this.addUndoElement(outputElement.innerHTML);
+
 				outputElement.style.lineHeight = this.markdownPreviewLineHeight.toString();
 				this.cellModel.renderedOutputTextContent = this.getRenderedTextOutput();
 				outputElement.focus();
@@ -262,6 +267,43 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 		this._changeRef.detectChanges();
 	}
 
+	private addUndoElement(newText: string) {
+		if (newText !== this._undoStack.peek()) {
+			this._redoStack.clear();
+			this._undoStack.push(newText);
+		}
+	}
+
+	private undoRichTextChange(): void {
+		// The first element in the undo stack is the initial cell text,
+		// which is the hard stop for undoing text changes. So we can only
+		// undo text changes after that one.
+		if (this._undoStack.count > 1) {
+			// The most recent change is at the top of the undo stack, so we want to
+			// update the text so that it's the change just before that.
+			let redoText = this._undoStack.pop();
+			this._redoStack.push(redoText);
+			let undoText = this._undoStack.peek();
+
+			let textOutputElement = <HTMLElement>this.output.nativeElement;
+			textOutputElement.innerHTML = undoText;
+
+			this.updateCellSource();
+		}
+	}
+
+	private redoRichTextChange(): void {
+		if (this._redoStack.count > 0) {
+			let text = this._redoStack.pop();
+			this._undoStack.push(text);
+
+			let textOutputElement = <HTMLElement>this.output.nativeElement;
+			textOutputElement.innerHTML = text;
+
+			this.updateCellSource();
+		}
+	}
+
 	//Sanitizes the content based on trusted mode of Cell Model
 	private sanitizeContent(content: string): string {
 		if (this.cellModel && !this.cellModel.trustedMode) {
@@ -286,6 +328,9 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 	}
 
 	public handleHtmlChanged(): void {
+		let textOutputElement = <HTMLElement>this.output.nativeElement;
+		this.addUndoElement(textOutputElement.innerHTML);
+
 		this.updateCellSource();
 	}
 
@@ -499,3 +544,50 @@ function preventDefaultAndExecCommand(e: KeyboardEvent, commandId: string) {
 	e.preventDefault();
 	document.execCommand(commandId);
 }
+
+/**
+ * A string stack used to track changes to Undo and Redo for the Rich Text editor in text cells.
+ */
+export class RichTextEditStack {
+	private _list: string[] = [];
+
+	/**
+	 * Adds an element to the top of the stack.
+	 * @param element The string element to add to the stack.
+	 */
+	public push(element: string): void {
+		this._list.push(element);
+	}
+
+	/**
+	 * Removes the topmost element of the stack and returns it.
+	 */
+	public pop(): string | undefined {
+		return this._list.pop();
+	}
+
+	/**
+	 * Returns the topmost element of the stack without removing it.
+	 */
+	public peek(): string | undefined {
+		if (this._list.length > 0) {
+			return this._list[this._list.length - 1];
+		} else {
+			return undefined;
+		}
+	}
+
+	/**
+	 * Removes all elements from the stack.
+	 */
+	public clear(): void {
+		this._list = [];
+	}
+
+	/**
+	 * Returns the number of elements in the stack.
+	 */
+	public get count(): number {
+		return this._list.length;
+	}
+}

src/sql/workbench/contrib/notebook/test/electron-browser/notebookUtils.test.ts

@@ -11,6 +11,7 @@ import { getHostAndPortFromEndpoint, isStream, getProvidersForFileName, asyncFor
 import { INotebookService, DEFAULT_NOTEBOOK_FILETYPE, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/browser/notebookService';
 import { NotebookServiceStub } from 'sql/workbench/contrib/notebook/test/stubs';
 import { tryMatchCellMagic, extractCellMagicCommandPlusArgs } from 'sql/workbench/services/notebook/browser/utils';
+import { RichTextEditStack } from 'sql/workbench/contrib/notebook/browser/cellViews/textCell.component';
 
 suite('notebookUtils', function (): void {
 	const mockNotebookService = TypeMoq.Mock.ofType<INotebookService>(NotebookServiceStub);
@@ -261,4 +262,33 @@ suite('notebookUtils', function (): void {
 		result = rewriteUrlUsingRegex(/(https?:\/\/sparkhead.*\/proxy)(.*)/g, html, '1.1.1.1', ':999', '/gateway/default/yarn/proxy');
 		assert.strictEqual(result, '<a target="_blank" href="https://storage-0-0.storage-0-svc.mssql-cluster.svc.cluster.local:8044/node/containerlogs/container_7/root“>Link</a>', 'Target URL should not have been edited');
 	});
+
+	test('EditStack test', async function (): Promise<void> {
+		let stack = new RichTextEditStack();
+		assert.strictEqual(stack.count, 0);
+
+		stack.push('1');
+		stack.push('2');
+		stack.push('3');
+		assert.strictEqual(stack.count, 3);
+
+		assert.strictEqual(stack.peek(), '3');
+
+		let topElement = stack.pop();
+		assert.strictEqual(topElement, '3');
+
+		topElement = stack.pop();
+		assert.strictEqual(topElement, '2');
+
+		stack.push('4');
+		assert.strictEqual(stack.count, 2);
+		topElement = stack.pop();
+		assert.strictEqual(topElement, '4');
+
+		stack.clear();
+		assert.strictEqual(stack.count, 0);
+		topElement = stack.pop();
+		assert.strictEqual(topElement, undefined);
+		assert.strictEqual(stack.peek(), undefined);
+	});
 });