From 769615708412f52612dc203366949508ab27eefc Mon Sep 17 00:00:00 2001
From: Cory Rivera <corivera@microsoft.com>
Date: Fri, 10 Jun 2022 17:15:50 -0700
Subject: [PATCH] Refresh cell toolbar when changing cell type. (#19704)

---
 .../cellViews/cellToolbar.component.ts        | 40 ++++++++++++++++---
 1 file changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/sql/workbench/contrib/notebook/browser/cellViews/cellToolbar.component.ts b/src/sql/workbench/contrib/notebook/browser/cellViews/cellToolbar.component.ts
index d8007cb38c..76664f5d05 100644
--- a/src/sql/workbench/contrib/notebook/browser/cellViews/cellToolbar.component.ts
+++ b/src/sql/workbench/contrib/notebook/browser/cellViews/cellToolbar.component.ts
@@ -17,6 +17,7 @@ import { DropdownMenuActionViewItem } from 'sql/base/browser/ui/buttonMenu/butto
 import { ICellModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/notebookModel';
 import { CellContext } from 'sql/workbench/contrib/notebook/browser/cellViews/codeActions';
+import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 
 export const CELL_TOOLBAR_SELECTOR: string = 'cell-toolbar-component';
 
@@ -39,7 +40,10 @@ export class CellToolbarComponent {
 	@Input() model: NotebookModel;
 
 	private _actionBar: Taskbar;
+	private _disposableActions: DisposableStore;
 	private _editCellAction: EditCellAction;
+	private _cellContext: CellContext;
+	private _typeChangedListener: IDisposable;
 	public _cellToggleMoreActions: CellToggleMoreActions;
 
 	constructor(
@@ -47,41 +51,65 @@ export class CellToolbarComponent {
 		@Inject(IContextMenuService) private contextMenuService: IContextMenuService
 	) {
 		this._cellToggleMoreActions = this.instantiationService.createInstance(CellToggleMoreActions);
+		this._disposableActions = new DisposableStore();
 	}
 
 	ngOnInit() {
 		this.initActionBar();
+		this._typeChangedListener = this.model.onCellTypeChanged(cell => {
+			if (cell === this.cellModel) {
+				this.setupActions();
+			}
+		});
+	}
+
+	ngOnDestroy() {
+		this._typeChangedListener.dispose();
 	}
 
 	protected initActionBar(): void {
-		let context = new CellContext(this.model, this.cellModel);
+		this._cellContext = new CellContext(this.model, this.cellModel);
 		let taskbar = <HTMLElement>this.celltoolbar.nativeElement;
 		this._actionBar = new Taskbar(taskbar);
-		this._actionBar.context = context;
+		this._actionBar.context = this._cellContext;
+
+		this.setupActions();
+	}
+
+	private setupActions(): void {
+		this._disposableActions.clear();
 
 		let addCellsButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('codeCellsPreview', "Add cell"), 'masked-pseudo code');
+		this._disposableActions.add(addCellsButton);
 
 		let addCodeCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddCodeCell', localize('codePreview', "Code cell"), 'masked-pseudo code');
 		addCodeCellButton.cellType = CellTypes.Code;
+		this._disposableActions.add(addCodeCellButton);
 
 		let addTextCellButton = this.instantiationService.createInstance(AddCellAction, 'notebook.AddTextCell', localize('textPreview', "Text cell"), 'masked-pseudo markdown');
 		addTextCellButton.cellType = CellTypes.Markdown;
+		this._disposableActions.add(addTextCellButton);
 
 		let moveCellDownButton = this.instantiationService.createInstance(MoveCellAction, 'notebook.MoveCellDown', 'masked-icon move-down', this.buttonMoveDown);
 		let moveCellUpButton = this.instantiationService.createInstance(MoveCellAction, 'notebook.MoveCellUp', 'masked-icon move-up', this.buttonMoveUp);
+		this._disposableActions.add(moveCellDownButton);
+		this._disposableActions.add(moveCellUpButton);
 
 		let splitCellButton = this.instantiationService.createInstance(SplitCellAction, 'notebook.SplitCellAtCursor', this.buttonSplitCell, 'masked-icon icon-split-cell');
-		splitCellButton.setListener(context);
+		splitCellButton.setListener(this._cellContext);
 		splitCellButton.enabled = this.cellModel.cellType !== 'markdown';
+		this._disposableActions.add(splitCellButton);
 
 		let deleteButton = this.instantiationService.createInstance(DeleteCellAction, 'notebook.DeleteCell', 'masked-icon delete', this.buttonDelete);
+		this._disposableActions.add(deleteButton);
 
 		let moreActionsContainer = DOM.$('li.action-item');
 		this._cellToggleMoreActions = this.instantiationService.createInstance(CellToggleMoreActions);
-		this._cellToggleMoreActions.onInit(moreActionsContainer, context);
+		this._cellToggleMoreActions.onInit(moreActionsContainer, this._cellContext);
 
 		this._editCellAction = this.instantiationService.createInstance(EditCellAction, 'notebook.EditCell', true, this.cellModel.isEditMode);
 		this._editCellAction.enabled = true;
+		this._disposableActions.add(this._editCellAction);
 
 		let addCellDropdownContainer = DOM.$('li.action-item');
 		addCellDropdownContainer.setAttribute('role', 'presentation');
@@ -97,10 +125,10 @@ export class CellToolbarComponent {
 			undefined
 		);
 		dropdownMenuActionViewItem.render(addCellDropdownContainer);
-		dropdownMenuActionViewItem.setActionContext(context);
+		dropdownMenuActionViewItem.setActionContext(this._cellContext);
 
 		let taskbarContent: ITaskbarContent[] = [];
-		if (this.cellModel?.cellType === CellTypes.Markdown) {
+		if (this.cellModel.cellType === CellTypes.Markdown) {
 			taskbarContent.push(
 				{ action: this._editCellAction }
 			);