Merge pull request #149195 from microsoft/roblourens/issue147527

Set correct editor selection when moving cursor from one cell editor to another. 

10 files changed, 62 insertions, 48 deletions

diff --git a/src/vs/workbench/contrib/notebook/browser/contrib/navigation/arrow.ts b/src/vs/workbench/contrib/notebook/browser/contrib/navigation/arrow.ts
index 763c427e321..b0dd1d172c8 100644
--- a/src/vs/workbench/contrib/notebook/browser/contrib/navigation/arrow.ts
+++ b/src/vs/workbench/contrib/notebook/browser/contrib/navigation/arrow.ts
@@ -32,7 +32,7 @@ const NOTEBOOK_CURSOR_PAGEDOWN_COMMAND_ID = 'notebook.cell.cursorPageDown';
 const NOTEBOOK_CURSOR_PAGEDOWN_SELECT_COMMAND_ID = 'notebook.cell.cursorPageDownSelect';
 
 
-registerAction2(class extends NotebookCellAction {
+registerAction2(class FocusNextCellAction extends NotebookCellAction {
 	constructor() {
 		super({
 			id: NOTEBOOK_FOCUS_NEXT_EDITOR,
@@ -76,13 +76,13 @@ registerAction2(class extends NotebookCellAction {
 
 		const newCell = editor.cellAt(idx + 1);
 		const newFocusMode = newCell.cellKind === CellKind.Markup && newCell.getEditState() === CellEditState.Preview ? 'container' : 'editor';
-		editor.focusNotebookCell(newCell, newFocusMode);
+		await editor.focusNotebookCell(newCell, newFocusMode, { focusEditorLine: 1 });
 		editor.cursorNavigationMode = true;
 	}
 });
 
 
-registerAction2(class extends NotebookCellAction {
+registerAction2(class FocusPreviousCellAction extends NotebookCellAction {
 	constructor() {
 		super({
 			id: NOTEBOOK_FOCUS_PREVIOUS_EDITOR,
@@ -118,7 +118,7 @@ registerAction2(class extends NotebookCellAction {
 
 		const newCell = editor.cellAt(idx - 1);
 		const newFocusMode = newCell.cellKind === CellKind.Markup && newCell.getEditState() === CellEditState.Preview ? 'container' : 'editor';
-		editor.focusNotebookCell(newCell, newFocusMode);
+		await editor.focusNotebookCell(newCell, newFocusMode, { focusEditorLine: newCell.textBuffer.getLineCount() });
 		editor.cursorNavigationMode = true;
 	}
 });
@@ -145,7 +145,7 @@ registerAction2(class extends NotebookAction {
 		}
 
 		const firstCell = editor.cellAt(0);
-		editor.focusNotebookCell(firstCell, 'container');
+		await editor.focusNotebookCell(firstCell, 'container');
 	}
 });
 
@@ -173,7 +173,7 @@ registerAction2(class extends NotebookAction {
 		const lastVisibleIdx = editor.getPreviousVisibleCellIndex(lastIdx);
 		if (lastVisibleIdx) {
 			const cell = editor.cellAt(lastVisibleIdx);
-			editor.focusNotebookCell(cell, 'container');
+			await editor.focusNotebookCell(cell, 'container');
 		}
 	}
 });
@@ -196,7 +196,7 @@ registerAction2(class extends NotebookCellAction {
 	async runWithContext(accessor: ServicesAccessor, context: INotebookCellActionContext): Promise<void> {
 		const editor = context.notebookEditor;
 		const activeCell = context.cell;
-		editor.focusNotebookCell(activeCell, 'output');
+		await editor.focusNotebookCell(activeCell, 'output');
 	}
 });
 
@@ -217,7 +217,7 @@ registerAction2(class extends NotebookCellAction {
 	async runWithContext(accessor: ServicesAccessor, context: INotebookCellActionContext): Promise<void> {
 		const editor = context.notebookEditor;
 		const activeCell = context.cell;
-		editor.focusNotebookCell(activeCell, 'editor');
+		await editor.focusNotebookCell(activeCell, 'editor');
 	}
 });
 
diff --git a/src/vs/workbench/contrib/notebook/browser/controller/cellOperations.ts b/src/vs/workbench/contrib/notebook/browser/controller/cellOperations.ts
index f4209f52ad9..08f3f326881 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/cellOperations.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/cellOperations.ts
@@ -69,7 +69,7 @@ export async function changeCellToKind(kind: CellKind, context: INotebookActionC
 			};
 		}, undefined, true);
 		const newCell = notebookEditor.cellAt(idx);
-		notebookEditor.focusNotebookCell(newCell, cell.getEditState() === CellEditState.Editing ? 'editor' : 'container');
+		await notebookEditor.focusNotebookCell(newCell, cell.getEditState() === CellEditState.Editing ? 'editor' : 'container');
 	} else if (context.selectedCells) {
 		const selectedCells = context.selectedCells;
 		const rawEdits: ICellEditOperation[] = [];
diff --git a/src/vs/workbench/contrib/notebook/browser/controller/editActions.ts b/src/vs/workbench/contrib/notebook/browser/controller/editActions.ts
index 8d253828d2e..7250d2dc560 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/editActions.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/editActions.ts
@@ -87,7 +87,7 @@ registerAction2(class EditCellAction extends NotebookCellAction {
 			return;
 		}
 
-		context.notebookEditor.focusNotebookCell(context.cell, 'editor');
+		await context.notebookEditor.focusNotebookCell(context.cell, 'editor');
 	}
 });
 
@@ -139,7 +139,7 @@ registerAction2(class QuitEditCellAction extends NotebookCellAction {
 			context.cell.updateEditState(CellEditState.Preview, QUIT_EDIT_CELL_COMMAND_ID);
 		}
 
-		context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
+		await context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
 	}
 });
 
@@ -499,7 +499,7 @@ async function setCellToLanguage(languageId: string, context: IChangeCellContext
 		const newCell = context.notebookEditor.cellAt(idx);
 
 		if (newCell) {
-			context.notebookEditor.focusNotebookCell(newCell, 'editor');
+			await context.notebookEditor.focusNotebookCell(newCell, 'editor');
 		}
 	} else if (languageId !== 'markdown' && context.cell?.cellKind === CellKind.Markup) {
 		await changeCellToKind(CellKind.Code, { cell: context.cell, notebookEditor: context.notebookEditor, ui: true }, languageId);
diff --git a/src/vs/workbench/contrib/notebook/browser/controller/executeActions.ts b/src/vs/workbench/contrib/notebook/browser/controller/executeActions.ts
index 063d3812d0e..564cd6548e5 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/executeActions.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/executeActions.ts
@@ -194,7 +194,7 @@ registerAction2(class ExecuteCell extends NotebookMultiCellAction {
 
 	async runWithContext(accessor: ServicesAccessor, context: INotebookCommandContext | INotebookCellToolbarActionContext): Promise<void> {
 		if (context.ui) {
-			context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
+			await context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
 		}
 
 		return runCell(accessor, context);
@@ -235,7 +235,7 @@ registerAction2(class ExecuteAboveCells extends NotebookMultiCellAction {
 		let endCellIdx: number | undefined = undefined;
 		if (context.ui) {
 			endCellIdx = context.notebookEditor.getCellIndex(context.cell);
-			context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
+			await context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
 		} else {
 			endCellIdx = Math.min(...context.selectedCells.map(cell => context.notebookEditor.getCellIndex(cell)));
 		}
@@ -282,7 +282,7 @@ registerAction2(class ExecuteCellAndBelow extends NotebookMultiCellAction {
 		let startCellIdx: number | undefined = undefined;
 		if (context.ui) {
 			startCellIdx = context.notebookEditor.getCellIndex(context.cell);
-			context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
+			await context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
 		} else {
 			startCellIdx = Math.min(...context.selectedCells.map(cell => context.notebookEditor.getCellIndex(cell)));
 		}
@@ -315,12 +315,12 @@ registerAction2(class ExecuteCellFocusContainer extends NotebookMultiCellAction
 
 	async runWithContext(accessor: ServicesAccessor, context: INotebookCommandContext | INotebookCellToolbarActionContext): Promise<void> {
 		if (context.ui) {
-			context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
+			await context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
 		} else {
 			const firstCell = context.selectedCells[0];
 
 			if (firstCell) {
-				context.notebookEditor.focusNotebookCell(firstCell, 'container', { skipReveal: true });
+				await context.notebookEditor.focusNotebookCell(firstCell, 'container', { skipReveal: true });
 			}
 		}
 
@@ -390,7 +390,7 @@ registerAction2(class CancelExecuteCell extends NotebookMultiCellAction {
 
 	async runWithContext(accessor: ServicesAccessor, context: INotebookCommandContext | INotebookCellToolbarActionContext): Promise<void> {
 		if (context.ui) {
-			context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
+			await context.notebookEditor.focusNotebookCell(context.cell, 'container', { skipReveal: true });
 			return context.notebookEditor.cancelNotebookCells(Iterable.single(context.cell));
 		} else {
 			return context.notebookEditor.cancelNotebookCells(context.selectedCells);
@@ -423,12 +423,12 @@ registerAction2(class ExecuteCellSelectBelow extends NotebookCellAction {
 			const nextCell = context.notebookEditor.cellAt(idx + 1);
 			context.cell.updateEditState(CellEditState.Preview, EXECUTE_CELL_SELECT_BELOW);
 			if (nextCell) {
-				context.notebookEditor.focusNotebookCell(nextCell, 'container');
+				await context.notebookEditor.focusNotebookCell(nextCell, 'container');
 			} else {
 				const newCell = insertCell(languageService, context.notebookEditor, idx, CellKind.Markup, 'below');
 
 				if (newCell) {
-					context.notebookEditor.focusNotebookCell(newCell, 'editor');
+					await context.notebookEditor.focusNotebookCell(newCell, 'editor');
 				}
 			}
 			return;
@@ -436,12 +436,12 @@ registerAction2(class ExecuteCellSelectBelow extends NotebookCellAction {
 			// Try to select below, fall back on inserting
 			const nextCell = context.notebookEditor.cellAt(idx + 1);
 			if (nextCell) {
-				context.notebookEditor.focusNotebookCell(nextCell, 'container');
+				await context.notebookEditor.focusNotebookCell(nextCell, 'container');
 			} else {
 				const newCell = insertCell(languageService, context.notebookEditor, idx, CellKind.Code, 'below');
 
 				if (newCell) {
-					context.notebookEditor.focusNotebookCell(newCell, 'editor');
+					await context.notebookEditor.focusNotebookCell(newCell, 'editor');
 				}
 			}
 
@@ -471,7 +471,7 @@ registerAction2(class ExecuteCellInsertBelow extends NotebookCellAction {
 
 		const newCell = insertCell(languageService, context.notebookEditor, idx, context.cell.cellKind, 'below');
 		if (newCell) {
-			context.notebookEditor.focusNotebookCell(newCell, newFocusMode);
+			await context.notebookEditor.focusNotebookCell(newCell, newFocusMode);
 		}
 
 		if (context.cell.cellKind === CellKind.Markup) {
diff --git a/src/vs/workbench/contrib/notebook/browser/controller/insertCellActions.ts b/src/vs/workbench/contrib/notebook/browser/controller/insertCellActions.ts
index 79d40142457..6a85c41a24a 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/insertCellActions.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/insertCellActions.ts
@@ -54,7 +54,7 @@ abstract class InsertCellCommand extends NotebookAction {
 		}
 
 		if (newCell) {
-			context.notebookEditor.focusNotebookCell(newCell, this.focusEditor ? 'editor' : 'container');
+			await context.notebookEditor.focusNotebookCell(newCell, this.focusEditor ? 'editor' : 'container');
 		}
 	}
 }
@@ -189,7 +189,7 @@ registerAction2(class InsertCodeCellAtTopAction extends NotebookAction {
 		const newCell = insertCell(languageService, context.notebookEditor, 0, CellKind.Code, 'above', undefined, true);
 
 		if (newCell) {
-			context.notebookEditor.focusNotebookCell(newCell, 'editor');
+			await context.notebookEditor.focusNotebookCell(newCell, 'editor');
 		}
 	}
 });
@@ -216,7 +216,7 @@ registerAction2(class InsertMarkdownCellAtTopAction extends NotebookAction {
 		const newCell = insertCell(languageService, context.notebookEditor, 0, CellKind.Markup, 'above', undefined, true);
 
 		if (newCell) {
-			context.notebookEditor.focusNotebookCell(newCell, 'editor');
+			await context.notebookEditor.focusNotebookCell(newCell, 'editor');
 		}
 	}
 });
diff --git a/src/vs/workbench/contrib/notebook/browser/diff/notebookDiffEditorBrowser.ts b/src/vs/workbench/contrib/notebook/browser/diff/notebookDiffEditorBrowser.ts
index 9bb383e2b81..26c6588808b 100644
--- a/src/vs/workbench/contrib/notebook/browser/diff/notebookDiffEditorBrowser.ts
+++ b/src/vs/workbench/contrib/notebook/browser/diff/notebookDiffEditorBrowser.ts
@@ -43,8 +43,8 @@ export interface INotebookTextDiffEditor {
 	 */
 	triggerScroll(event: IMouseWheelEvent): void;
 	getCellByInfo(cellInfo: ICommonCellInfo): IGenericCellViewModel;
-	focusNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output'): void;
-	focusNextNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output'): void;
+	focusNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output'): Promise<void>;
+	focusNextNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output'): Promise<void>;
 	updateOutputHeight(cellInfo: ICommonCellInfo, output: ICellOutputViewModel, height: number, isInit: boolean): void;
 	deltaCellOutputContainerClassNames(diffSide: DiffSide, cellId: string, added: string[], removed: string[]): void;
 }
diff --git a/src/vs/workbench/contrib/notebook/browser/diff/notebookTextDiffEditor.ts b/src/vs/workbench/contrib/notebook/browser/diff/notebookTextDiffEditor.ts
index 9de089e2f5a..1369de699b9 100644
--- a/src/vs/workbench/contrib/notebook/browser/diff/notebookTextDiffEditor.ts
+++ b/src/vs/workbench/contrib/notebook/browser/diff/notebookTextDiffEditor.ts
@@ -158,11 +158,11 @@ export class NotebookTextDiffEditor extends EditorPane implements INotebookTextD
 		// throw new Error('Method not implemented.');
 	}
 
-	focusNotebookCell(cell: IGenericCellViewModel, focus: 'output' | 'editor' | 'container'): void {
+	async focusNotebookCell(cell: IGenericCellViewModel, focus: 'output' | 'editor' | 'container'): Promise<void> {
 		// throw new Error('Method not implemented.');
 	}
 
-	focusNextNotebookCell(cell: IGenericCellViewModel, focus: 'output' | 'editor' | 'container'): void {
+	async focusNextNotebookCell(cell: IGenericCellViewModel, focus: 'output' | 'editor' | 'container'): Promise<void> {
 		// throw new Error('Method not implemented.');
 	}
 
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts b/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts
index bac71426b9a..56012a1a041 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts
@@ -122,6 +122,7 @@ export interface ICommonCellInfo {
 
 export interface IFocusNotebookCellOptions {
 	readonly skipReveal?: boolean;
+	readonly focusEditorLine?: number;
 }
 
 //#endregion
@@ -446,7 +447,7 @@ export interface INotebookEditor {
 	/**
 	 * Focus the container of a cell (the monaco editor inside is not focused).
 	 */
-	focusNotebookCell(cell: ICellViewModel, focus: 'editor' | 'container' | 'output', options?: IFocusNotebookCellOptions): void;
+	focusNotebookCell(cell: ICellViewModel, focus: 'editor' | 'container' | 'output', options?: IFocusNotebookCellOptions): Promise<void>;
 
 	/**
 	 * Execute the given notebook cells
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
index 69aa0a067d8..535e00b2aa6 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
@@ -1077,14 +1077,14 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		}));
 	}
 
-	private _updateForCursorNavigationMode(applyFocusChange: () => void): void {
+	private async _updateForCursorNavigationMode(applyFocusChange: () => void): Promise<void> {
 		if (this._cursorNavigationMode) {
 			// Will fire onDidChangeFocus, resetting the state to Container
 			applyFocusChange();
 
 			const newFocusedCell = this._list.getFocusedElements()[0];
 			if (newFocusedCell.cellKind === CellKind.Code || newFocusedCell.getEditState() === CellEditState.Editing) {
-				this.focusNotebookCell(newFocusedCell, 'editor');
+				await this.focusNotebookCell(newFocusedCell, 'editor');
 			} else {
 				// Reset to "Editor", the state has not been consumed
 				this._cursorNavigationMode = true;
@@ -2326,7 +2326,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		return selectedCellsInRange;
 	}
 
-	focusNotebookCell(cell: ICellViewModel, focusItem: 'editor' | 'container' | 'output', options?: IFocusNotebookCellOptions) {
+	async focusNotebookCell(cell: ICellViewModel, focusItem: 'editor' | 'container' | 'output', options?: IFocusNotebookCellOptions) {
 		if (this._isDisposed) {
 			return;
 		}
@@ -2339,13 +2339,26 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 			cell.updateEditState(CellEditState.Editing, 'focusNotebookCell');
 			cell.focusMode = CellFocusMode.Editor;
 			if (!options?.skipReveal) {
-				const selectionsStartPosition = cell.getSelectionsStartPosition();
-				if (selectionsStartPosition?.length) {
-					const firstSelectionPosition = selectionsStartPosition[0];
-					this.revealRangeInCenterIfOutsideViewportAsync(cell, Range.fromPositions(firstSelectionPosition, firstSelectionPosition));
+				if (typeof options?.focusEditorLine === 'number') {
+					await this.revealLineInViewAsync(cell, options.focusEditorLine);
+					const editor = this._renderedEditors.get(cell)!;
+					const focusEditorLine = options.focusEditorLine!;
+					editor?.setSelection({
+						startLineNumber: focusEditorLine,
+						startColumn: 1,
+						endLineNumber: focusEditorLine,
+						endColumn: 1
+					});
 				} else {
-					this.revealInCenterIfOutsideViewport(cell);
+					const selectionsStartPosition = cell.getSelectionsStartPosition();
+					if (selectionsStartPosition?.length) {
+						const firstSelectionPosition = selectionsStartPosition[0];
+						await this.revealRangeInCenterIfOutsideViewportAsync(cell, Range.fromPositions(firstSelectionPosition, firstSelectionPosition));
+					} else {
+						this.revealInCenterIfOutsideViewport(cell);
+					}
 				}
+
 			}
 		} else if (focusItem === 'output') {
 			this.focusElement(cell);
@@ -2384,7 +2397,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		}
 	}
 
-	focusNextNotebookCell(cell: ICellViewModel, focusItem: 'editor' | 'container' | 'output') {
+	async focusNextNotebookCell(cell: ICellViewModel, focusItem: 'editor' | 'container' | 'output') {
 		const idx = this.viewModel?.getCellIndex(cell);
 		if (typeof idx !== 'number') {
 			return;
@@ -2395,7 +2408,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 			return;
 		}
 
-		this.focusNotebookCell(newCell, focusItem);
+		await this.focusNotebookCell(newCell, focusItem);
 	}
 
 	//#endregion
diff --git a/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts b/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts
index 70abc4be530..cc2eb8bf93a 100644
--- a/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts
+++ b/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts
@@ -63,10 +63,10 @@ export interface IResolvedBackLayerWebview {
 export interface INotebookDelegateForWebview {
 	readonly creationOptions: INotebookEditorCreationOptions;
 	getCellById(cellId: string): IGenericCellViewModel | undefined;
-	focusNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output', options?: IFocusNotebookCellOptions): void;
+	focusNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output', options?: IFocusNotebookCellOptions): Promise<void>;
 	toggleNotebookCellSelection(cell: IGenericCellViewModel, selectFromPrevious: boolean): void;
 	getCellByInfo(cellInfo: ICommonCellInfo): IGenericCellViewModel;
-	focusNextNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output'): void;
+	focusNextNotebookCell(cell: IGenericCellViewModel, focus: 'editor' | 'container' | 'output'): Promise<void>;
 	updateOutputHeight(cellInfo: ICommonCellInfo, output: IDisplayOutputViewModel, height: number, isInit: boolean, source?: string): void;
 	scheduleOutputHeightAck(cellInfo: ICommonCellInfo, outputId: string, height: number): void;
 	updateMarkupCellHeight(cellId: string, height: number, isInit: boolean): void;
@@ -560,7 +560,7 @@ var requirejs = (function() {
 			}
 		}));
 
-		this._register(this.webview.onMessage((message) => {
+		this._register(this.webview.onMessage(async (message) => {
 			const data: FromWebviewMessage | { readonly __vscode_notebook_message: undefined } = message.message;
 			if (this._disposed) {
 				return;
@@ -656,7 +656,7 @@ var requirejs = (function() {
 						if (data.focusNext) {
 							this.notebookEditor.focusNextNotebookCell(cell, 'editor');
 						} else {
-							this.notebookEditor.focusNotebookCell(cell, 'editor');
+							await this.notebookEditor.focusNotebookCell(cell, 'editor');
 						}
 					}
 					break;
@@ -734,7 +734,7 @@ var requirejs = (function() {
 							this.notebookEditor.toggleNotebookCellSelection(cell, /* fromPrevious */ data.shiftKey);
 						} else {
 							// Normal click
-							this.notebookEditor.focusNotebookCell(cell, 'container', { skipReveal: true });
+							await this.notebookEditor.focusNotebookCell(cell, 'container', { skipReveal: true });
 						}
 					}
 					break;
@@ -743,7 +743,7 @@ var requirejs = (function() {
 					const cell = this.notebookEditor.getCellById(data.cellId);
 					if (cell) {
 						// Focus the cell first
-						this.notebookEditor.focusNotebookCell(cell, 'container', { skipReveal: true });
+						await this.notebookEditor.focusNotebookCell(cell, 'container', { skipReveal: true });
 
 						// Then show the context menu
 						const webviewRect = this.element.getBoundingClientRect();
@@ -767,7 +767,7 @@ var requirejs = (function() {
 					const cell = this.notebookEditor.getCellById(data.cellId);
 					if (cell && !this.notebookEditor.creationOptions.isReadOnly) {
 						this.notebookEditor.setMarkupCellEditState(data.cellId, CellEditState.Editing);
-						this.notebookEditor.focusNotebookCell(cell, 'editor', { skipReveal: true });
+						await this.notebookEditor.focusNotebookCell(cell, 'editor', { skipReveal: true });
 					}
 					break;
 				}