Fix find highlight on cell edits (#22080)

* update the rendered text onCellSourceChange

* fix test

* fix highlight in split mode

* update corresponding test

* update hasEditor with getEditor

* update event

* add comment

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

@@ -337,6 +337,7 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 				if (this.isFindActive) {
 					this.addDecoration();
 				}
+				this.cellModel.cellPreviewUpdated.fire();
 			}
 		}
 	}

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

@@ -232,7 +232,10 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 				let cells = [...new Set(oldDecorationsRange.map(item => item.cell))].filter(c => c.cellType === 'markdown');
 				cells.forEach(cell => {
 					let cellOldDecorations = oldDecorationsRange.filter(r => r.cell === cell);
-					let cellEditor = this.cellEditors.find(c => c.cellGuid() === cell.cellGuid);
+					// In split mode we have code and text cells with the same Guid,
+					// we need to find the text component editor to apply the decorations
+					// text component doesn't have an editor => !c.getEditor() filters that.
+					let cellEditor = this.cellEditors.find(c => c.cellGuid() === cell.cellGuid && !c.getEditor());
 					cellEditor.deltaDecorations(undefined, cellOldDecorations);
 				});
 				// code cell outputs
@@ -244,7 +247,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 				});
 			} else {
 				if (oldDecorationsRange.cell.cellType === 'markdown' || oldDecorationsRange.outputComponentIndex >= 0) {
-					let cell = oldDecorationsRange.outputComponentIndex >= 0 ? this.cellEditors.filter(c => c.cellGuid() === oldDecorationsRange.cell.cellGuid && c.isCellOutput)[oldDecorationsRange.outputComponentIndex] : this.cellEditors.find(c => c.cellGuid() === oldDecorationsRange.cell.cellGuid);
+					let cell = oldDecorationsRange.outputComponentIndex >= 0 ? this.cellEditors.filter(c => c.cellGuid() === oldDecorationsRange.cell.cellGuid && c.isCellOutput)[oldDecorationsRange.outputComponentIndex] : this.cellEditors.find(c => c.cellGuid() === oldDecorationsRange.cell.cellGuid && !c.getEditor());
 					cell.deltaDecorations(undefined, oldDecorationsRange);
 				}
 			}
@@ -254,7 +257,10 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 				let cells = [...new Set(newDecorationsRange.map(item => item.cell))].filter(c => c.cellType === 'markdown');
 				cells.forEach(cell => {
 					let cellNewDecorations = newDecorationsRange.filter(r => r.cell === cell);
-					let cellEditor = this.cellEditors.find(c => c.cellGuid() === cell.cellGuid);
+					// In split mode we have code and text cells with the same Guid,
+					// we need to find the text component editor to apply the decorations
+					// text component doesn't have an editor => !c.getEditor() filters that.
+					let cellEditor = this.cellEditors.find(c => c.cellGuid() === cell.cellGuid && !c.getEditor());
 					cellEditor.deltaDecorations(cellNewDecorations, undefined);
 				});
 				// code cell outputs
@@ -266,7 +272,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 				});
 			} else {
 				if (newDecorationsRange.cell.cellType === 'markdown' || newDecorationsRange.outputComponentIndex >= 0) {
-					let cell = newDecorationsRange.outputComponentIndex >= 0 ? this.cellEditors.filter(c => c.cellGuid() === newDecorationsRange.cell.cellGuid && c.isCellOutput)[newDecorationsRange.outputComponentIndex] : this.cellEditors.find(c => c.cellGuid() === newDecorationsRange.cell.cellGuid);
+					let cell = newDecorationsRange.outputComponentIndex >= 0 ? this.cellEditors.filter(c => c.cellGuid() === newDecorationsRange.cell.cellGuid && c.isCellOutput)[newDecorationsRange.outputComponentIndex] : this.cellEditors.find(c => c.cellGuid() === newDecorationsRange.cell.cellGuid && !c.getEditor());
 					cell.deltaDecorations(newDecorationsRange, undefined);
 				}
 			}

src/sql/workbench/contrib/notebook/browser/notebookEditor.ts

@@ -26,7 +26,7 @@ import { NotebookFindNextAction, NotebookFindPreviousAction } from 'sql/workbenc
 import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
 import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 import { IWorkbenchThemeService } from 'vs/workbench/services/themes/common/workbenchThemeService';
-import { CellEditModes, INotebookModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
+import { INotebookModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { INotebookFindModel } from 'sql/workbench/contrib/notebook/browser/models/notebookFindModel';
 import { IDisposable, DisposableStore } from 'vs/base/common/lifecycle';
 import { IModelDecorationsChangeAccessor, IModelDeltaDecoration } from 'vs/editor/common/model';
@@ -343,13 +343,10 @@ export class NotebookEditor extends EditorPane implements IFindNotebookControlle
 			}
 		}
 		if (e.searchScope) {
+			this._findDecorations.clearDecorations();
 			await this.notebookInput.notebookFindModel.find(this._findState.searchString, this._findState.matchCase, this._findState.wholeWord, NOTEBOOK_MAX_MATCHES);
 			this._findDecorations.set(this.notebookFindModel.findMatches, this.notebookFindModel.findArray);
-			this._findState.changeMatchInfo(
+			this._updateFinderMatchState();
-				this.notebookFindModel.getIndexByRange(this._currentMatch),
-				this._findDecorations.getCount(),
-				this._currentMatch
-			);
 		}
 	}
 
@@ -374,15 +371,11 @@ export class NotebookEditor extends EditorPane implements IFindNotebookControlle
 			filters: false
 		};
 		this._notebookModel.cells?.forEach(cell => {
-			this._register(cell.onCellEditModeChanged((isEditMode) => {
-				if (isEditMode) {
-					this._onFindStateChange(changeEvent).catch(onUnexpectedError);
-				}
-			}));
 			this._register(cell.onCurrentEditModeChanged(editMode => {
-				if (editMode !== CellEditModes.WYSIWYG) {
+				this._onFindStateChange(changeEvent).catch(onUnexpectedError);
-					this._onFindStateChange(changeEvent).catch(onUnexpectedError);
+			}));
-				}
+			this._register(cell.onCellPreviewUpdated(previewUpdated => {
+				this._onFindStateChange(changeEvent).catch(onUnexpectedError);
 			}));
 		});
 		this._register(this._notebookModel.contentChanged(e => {

src/sql/workbench/contrib/notebook/test/browser/notebookEditor.test.ts

@@ -15,7 +15,7 @@ import { NotebookEditor } from 'sql/workbench/contrib/notebook/browser/notebookE
 import { NBTestQueryManagementService } from 'sql/workbench/contrib/notebook/test/nbTestQueryManagementService';
 import * as stubs from 'sql/workbench/contrib/notebook/test/stubs';
 import { NotebookEditorStub } from 'sql/workbench/contrib/notebook/test/testCommon';
-import { ICellModel, NotebookContentChange } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
+import { CellEditModes, ICellModel, NotebookContentChange } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { INotebookEditor, INotebookParams, INotebookService, NotebookRange } from 'sql/workbench/services/notebook/browser/notebookService';
 import { NotebookService } from 'sql/workbench/services/notebook/browser/notebookServiceImpl';
 import { NotebookChangeType } from 'sql/workbench/services/notebook/common/contracts';
@@ -439,7 +439,7 @@ suite('Test class NotebookEditor:', () => {
 				TypeMoq.It.isAny(),
 				TypeMoq.It.isAnyNumber()
 			), TypeMoq.Times.once());
-			findDecorationsMock.verify(x => x.clearDecorations(), TypeMoq.Times.never());
+			findDecorationsMock.verify(x => x.clearDecorations(), TypeMoq.Times.once());
 			notebookFindModelMock.verify(x => x.clearFind(), TypeMoq.Times.never());
 			if (visibility === 'visible') {
 				assert.strictEqual(notebookEditor.getPosition(), currentMatch, `position must be set to the same NotebookRange that we set for '_currentMatch' property of notebookEditor`);
@@ -463,7 +463,7 @@ suite('Test class NotebookEditor:', () => {
 		const notebookModel = <NotebookModelStub>await notebookEditor.getNotebookModel();
 		notebookModel['_cells'] = [new CellModel({ cell_type: 'code', source: '' }, { isTrusted: true, notebook: notebookModel })];
 		notebookEditor['registerModelChanges']();
-		notebookModel.cells[0]['_onCellEditModeChanged'].fire(true); //fire cellEditModeChanged event on the first sell of our test notebookModel
+		notebookModel.cells[0]['_onCurrentEditModeChanged'].fire(CellEditModes.SPLIT); //fire onCurrentEditModeChanged event on the first cell of our test notebookModel
 		notebookModel.contentChangedEmitter.fire({ changeType: NotebookChangeType.Saved });
 		(<NotebookService>notebookService)['_onNotebookEditorAdd'].fire(<INotebookEditor>{});
 		notebookFindModelMock.verify(x => x.find(

src/sql/workbench/services/notebook/browser/models/cell.ts

@@ -92,6 +92,7 @@ export class CellModel extends Disposable implements ICellModel {
 	private _lastEditMode: string | undefined;
 	public richTextCursorPosition: ICaretPosition | undefined;
 	public markdownCursorPosition: IPosition | undefined;
+	public cellPreviewUpdated = new Emitter<void>();
 
 	constructor(cellData: nb.ICellContents,
 		private _options: ICellModelOptions,
@@ -509,6 +510,10 @@ export class CellModel extends Disposable implements ICellModel {
 		this._cellSourceChanged = val;
 	}
 
+	public get onCellPreviewUpdated(): Event<void> {
+		return this.cellPreviewUpdated.event;
+	}
+
 	public get onParameterStateChanged(): Event<boolean> {
 		return this._onParameterStateChanged.event;
 	}

src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts

@@ -6,7 +6,7 @@
 // This code is based on @jupyterlab/packages/apputils/src/clientsession.tsx
 
 import { nb } from 'azdata';
-import { Event } from 'vs/base/common/event';
+import { Emitter, Event } from 'vs/base/common/event';
 import { IDisposable } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { INotificationService } from 'vs/platform/notification/common/notification';
@@ -456,6 +456,8 @@ export interface ICellModel {
 	readonly onCollapseStateChanged: Event<boolean>;
 	readonly onParameterStateChanged: Event<boolean>;
 	readonly onCellEditModeChanged: Event<boolean>;
+	readonly onCellPreviewUpdated: Event<void>;
+	readonly cellPreviewUpdated: Emitter<void>;
 	modelContentChangedEvent: IModelContentChangedEvent;
 	isEditMode: boolean;
 	showPreview: boolean;