Keep selected kernel from untitled notebook when saving. (#152280)

Fixes #151158
Also fix cachine the previously selected kernel for Untitled-1 after saving it

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

@@ -331,6 +331,7 @@ export interface INotebookEditorViewState {
 	focus?: number;
 	editorFocused?: boolean;
 	contributionsState?: { [id: string]: unknown };
+	selectedKernelId?: string;
 }
 
 export interface ICellModelDecorations {

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

@@ -1120,6 +1120,8 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 			this.restoreListViewState(viewState);
 		}
 
+		this._restoreSelectedKernel(viewState);
+
 		// load preloads for matching kernel
 		this._loadKernelPreloads();
 
@@ -1685,6 +1687,15 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		}
 	}
 
+	private _restoreSelectedKernel(viewState: INotebookEditorViewState | undefined): void {
+		if (viewState?.selectedKernelId && this.textModel) {
+			const kernel = this.notebookKernelService.getMatchingKernel(this.textModel).all.find(k => k.id === viewState.selectedKernelId);
+			if (kernel) {
+				this.notebookKernelService.selectKernelForNotebook(kernel, this.textModel);
+			}
+		}
+	}
+
 	getEditorViewState(): INotebookEditorViewState {
 		const state = this.viewModel?.getEditorViewState();
 		if (!state) {
@@ -1730,8 +1741,9 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 				contributionsState[id] = contribution.saveViewState();
 			}
 		}
-
 		state.contributionsState = contributionsState;
+		state.selectedKernelId = this.activeKernel?.id;
+
 		return state;
 	}
 

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

@@ -115,7 +115,7 @@ export class NotebookKernelService extends Disposable implements INotebookKernel
 		this._register(_notebookService.onWillRemoveNotebookDocument(notebook => {
 			const kernelId = this._notebookBindings.get(NotebookTextModelLikeId.str(notebook));
 			if (kernelId) {
-				this._onDidChangeNotebookKernelBinding.fire({ notebook: notebook.uri, oldKernel: kernelId, newKernel: undefined });
+				this.selectKernelForNotebook(undefined, notebook);
 			}
 		}));
 		this._sourceMenu = this._register(this._menuService.createMenu(MenuId.NotebookKernelSource, contextKeyService));
@@ -287,7 +287,7 @@ export class NotebookKernelService extends Disposable implements INotebookKernel
 
 	// a notebook has one kernel, a kernel has N notebooks
 	// notebook <-1----N-> kernel
-	selectKernelForNotebook(kernel: INotebookKernel, notebook: INotebookTextModelLike): void {
+	selectKernelForNotebook(kernel: INotebookKernel | undefined, notebook: INotebookTextModelLike): void {
 		const key = NotebookTextModelLikeId.str(notebook);
 		const oldKernel = this._notebookBindings.get(key);
 		if (oldKernel !== kernel?.id) {
@@ -296,7 +296,7 @@ export class NotebookKernelService extends Disposable implements INotebookKernel
 			} else {
 				this._notebookBindings.delete(key);
 			}
-			this._onDidChangeNotebookKernelBinding.fire({ notebook: notebook.uri, oldKernel, newKernel: kernel.id });
+			this._onDidChangeNotebookKernelBinding.fire({ notebook: notebook.uri, oldKernel, newKernel: kernel?.id });
 			this._persistMementos();
 		}
 	}