handle custom context menu later.

3 files changed, 6 insertions, 17 deletions

diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
index aef15c3d580..c60dca406e8 100644
--- a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
+++ b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts
@@ -1544,15 +1544,14 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 
 	private _lastCellWithEditorFocus: ICellViewModel | null = null;
 	private _validateCellFocusMode(cell: ICellViewModel) {
-		if (cell.focusMode !== CellFocusMode.Editor) {
-			return;
-		}
+		if (cell.focusMode === CellFocusMode.Editor) {
+			if (this._lastCellWithEditorFocus && this._lastCellWithEditorFocus !== cell) {
+				this._lastCellWithEditorFocus.focusMode = CellFocusMode.Container;
+			}
 
-		if (this._lastCellWithEditorFocus && this._lastCellWithEditorFocus !== cell) {
-			this._lastCellWithEditorFocus.focusMode = CellFocusMode.Container;
+			this._lastCellWithEditorFocus = cell;
 		}
 
-		this._lastCellWithEditorFocus = cell;
 	}
 
 	private async _warmupWithMarkdownRenderer(viewModel: NotebookViewModel, viewState: INotebookEditorViewState | undefined) {
diff --git a/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts b/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts
index 515195332f5..63cc7be3072 100644
--- a/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts
+++ b/src/vs/workbench/contrib/notebook/browser/view/cellParts/codeCell.ts
@@ -330,11 +330,6 @@ export class CodeCell extends Disposable {
 			this.templateData.editor?.focus();
 		}
 
-		if (this.notebookEditor.getActiveCell() === this.viewCell
-			&& this.viewCell.focusMode === CellFocusMode.Container) {
-			this.notebookEditor.focusNotebookCell(this.viewCell, 'container');
-		}
-
 		this.templateData.container.classList.toggle('cell-editor-focus', this.viewCell.focusMode === CellFocusMode.Editor);
 		this.templateData.container.classList.toggle('cell-output-focus', this.viewCell.focusMode === CellFocusMode.Output);
 	}
diff --git a/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts b/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts
index 8f9fd7a1cac..5c9f559110d 100644
--- a/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts
+++ b/src/vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel.ts
@@ -12,7 +12,7 @@ import { ITextModelService } from 'vs/editor/common/services/resolverService';
 import { PrefixSumComputer } from 'vs/editor/common/model/prefixSumComputer';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
-import { CellEditState, CellFindMatch, CodeCellLayoutChangeEvent, CodeCellLayoutInfo, CellLayoutState, ICellOutputViewModel, ICellViewModel, CellFocusMode } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+import { CellEditState, CellFindMatch, CodeCellLayoutChangeEvent, CodeCellLayoutInfo, CellLayoutState, ICellOutputViewModel, ICellViewModel } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { CellOutputViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/cellOutputViewModel';
 import { ViewContext } from 'vs/workbench/contrib/notebook/browser/viewModel/viewContext';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
@@ -133,11 +133,6 @@ export class CodeCellViewModel extends BaseCellViewModel implements ICellViewMod
 			this._onDidRemoveOutputs.fire(removedOutputs);
 			this.layoutChange({ outputHeight: true }, 'CodeCellViewModel#model.onDidChangeOutputs');
 			dispose(removedOutputs);
-
-			if (this.outputsViewModels.length === 0 && this.focusMode === CellFocusMode.Output) {
-				// If the last output is removed, focus goes to container
-				this.focusMode = CellFocusMode.Container;
-			}
 		}));
 
 		this._outputCollection = new Array(this.model.outputs.length);