Persist focused notebook cell, dom focus on load

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

@@ -5,7 +5,7 @@
 
 import { IMouseWheelEvent } from 'vs/base/browser/mouseEvent';
 import { IListEvent, IListMouseEvent } from 'vs/base/browser/ui/list/list';
-import { IListOptions, IListStyles } from 'vs/base/browser/ui/list/listWidget';
+import { IListOptions, IListStyles, List } from 'vs/base/browser/ui/list/listWidget';
 import { ProgressBar } from 'vs/base/browser/ui/progressbar/progressbar';
 import { ToolBar } from 'vs/base/browser/ui/toolbar/toolbar';
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
@@ -341,6 +341,11 @@ export interface INotebookCellList {
 	updateOptions(options: IListOptions<ICellViewModel>): void;
 	layout(height?: number, width?: number): void;
 	dispose(): void;
+
+	// TODO resolve differences between List<CellViewModel> and INotebookCellList<ICellViewModel>
+	getFocus(): number[];
+	setFocus(indexes: number[]): void;
+	setSelection(indexes: number[]): void;
 }
 
 export interface BaseCellRenderTemplate {

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

@@ -312,6 +312,7 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 	focus() {
 		super.focus();
 		this.editorFocus?.set(true);
+		this.list?.domFocus();
 	}
 
 	async setInput(input: NotebookEditorInput, options: EditorOptions | undefined, token: CancellationToken): Promise<void> {
@@ -431,6 +432,10 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 
 		this.list!.layout();
 
+		this.restoreTextEditorViewState(viewState);
+	}
+
+	private restoreTextEditorViewState(viewState: INotebookEditorViewState | undefined): void {
 		if (viewState?.scrollPosition !== undefined) {
 			this.list!.scrollTop = viewState!.scrollPosition.top;
 			this.list!.scrollLeft = viewState!.scrollPosition.left;
@@ -438,6 +443,12 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 			this.list!.scrollTop = 0;
 			this.list!.scrollLeft = 0;
 		}
+
+		if (typeof viewState?.focus === 'number') {
+			this.list!.setFocus([viewState.focus]);
+		} else {
+			this.list!.setFocus([0]);
+		}
 	}
 
 	private saveTextEditorViewState(input: NotebookEditorInput): void {
@@ -456,6 +467,11 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 				}
 
 				state.cellTotalHeights = cellHeights;
+
+				const focus = this.list.getFocus()[0];
+				if (focus) {
+					state.focus = focus;
+				}
 			}
 
 			this.editorMemento.saveEditorState(this.group, input.resource, state);

src/vs/workbench/contrib/notebook/browser/view/notebookCellList.ts

@@ -6,7 +6,7 @@
 import * as DOM from 'vs/base/browser/dom';
 import { IMouseWheelEvent } from 'vs/base/browser/mouseEvent';
 import { IListRenderer, IListVirtualDelegate, ListError } from 'vs/base/browser/ui/list/list';
-import { IListStyles, IStyleController, MouseController } from 'vs/base/browser/ui/list/listWidget';
+import { IListStyles, IStyleController, MouseController, IListOptions } from 'vs/base/browser/ui/list/listWidget';
 import { Emitter, Event } from 'vs/base/common/event';
 import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { isMacintosh } from 'vs/base/common/platform';
@@ -120,7 +120,7 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 
 	}
 
-	protected createMouseController(): MouseController<CellViewModel> {
+	protected createMouseController(_options: IListOptions<CellViewModel>): MouseController<CellViewModel> {
 		return new NotebookMouseController(this);
 	}
 

src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts

@@ -33,6 +33,7 @@ export interface INotebookEditorViewState {
 	editorViewStates: { [key: number]: editorCommon.ICodeEditorViewState | null };
 	cellTotalHeights?: { [key: number]: number };
 	scrollPosition?: { left: number; top: number; };
+	focus?: number;
 }
 
 export interface ICellModelDecorations {
@@ -578,8 +579,8 @@ export class NotebookViewModel extends Disposable implements FoldingRegionDelega
 	}
 
 	saveEditorViewState(): INotebookEditorViewState {
-		const state: { [key: number]: boolean } = {};
-		this._viewCells.filter(cell => cell.editState === CellEditState.Editing).forEach(cell => state[cell.model.handle] = true);
+		const editingCells: { [key: number]: boolean } = {};
+		this._viewCells.filter(cell => cell.editState === CellEditState.Editing).forEach(cell => editingCells[cell.model.handle] = true);
 		const editorViewStates: { [key: number]: editorCommon.ICodeEditorViewState } = {};
 		this._viewCells.map(cell => ({ handle: cell.model.handle, state: cell.saveEditorViewState() })).forEach(viewState => {
 			if (viewState.state) {
@@ -588,8 +589,8 @@ export class NotebookViewModel extends Disposable implements FoldingRegionDelega
 		});
 
 		return {
-			editingCells: state,
-			editorViewStates: editorViewStates
+			editingCells,
+			editorViewStates
 		};
 	}