Fix cell widget layout update

src/vs/workbench/contrib/notebook/browser/media/notebookCellChat.css

@@ -6,7 +6,7 @@
 .monaco-workbench .notebookOverlay .cell-chat-part {
 	display: none;
 	color: inherit;
-	padding: 6px;
+	padding: 0 6px;
 	border-radius: 6px;
 	border: 1px solid var(--vscode-inlineChat-border);
 	background: var(--vscode-inlineChat-background);

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

@@ -863,7 +863,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		// chat
 		styleSheets.push(`
 		.monaco-workbench .notebookOverlay .cell-chat-part {
-			margin: 0 ${cellRightMargin}px 8px 6px;
+			margin: 0 ${cellRightMargin}px 6px 6px;
 		}
 		`);
 

src/vs/workbench/contrib/notebook/browser/view/cellParts/cellExecution.ts

@@ -78,7 +78,7 @@ export class CellExecutionPart extends CellContentPart {
 			DOM.hide(this._executionOrderLabel);
 		} else {
 			DOM.show(this._executionOrderLabel);
-			const top = element.layoutInfo.editorHeight - 22 + element.layoutInfo.statusBarHeight + element.layoutInfo.chatHeight;
+			const top = element.layoutInfo.editorHeight - 22 + element.layoutInfo.statusBarHeight;
 			this._executionOrderLabel.style.top = `${top}px`;
 		}
 	}

src/vs/workbench/contrib/notebook/browser/view/cellParts/cellFocusIndicator.ts

@@ -80,7 +80,6 @@ export class CellFocusIndicator extends CellContentPart {
 			this.bottom.domNode.style.transform = `translateY(${indicatorPostion.bottomIndicatorTop}px)`;
 			this.left.setHeight(indicatorPostion.verticalIndicatorHeight);
 			this.right.setHeight(indicatorPostion.verticalIndicatorHeight);
-			this.codeFocusIndicator.setTop(element.layoutInfo.chatHeight);
 			this.codeFocusIndicator.setHeight(indicatorPostion.verticalIndicatorHeight - this.getIndicatorTopMargin() * 2 - element.layoutInfo.chatHeight);
 		} else {
 			const cell = element as CodeCellViewModel;
@@ -89,9 +88,7 @@ export class CellFocusIndicator extends CellContentPart {
 			const indicatorHeight = cell.layoutInfo.codeIndicatorHeight + cell.layoutInfo.outputIndicatorHeight + cell.layoutInfo.commentHeight;
 			this.left.setHeight(indicatorHeight);
 			this.right.setHeight(indicatorHeight);
-			this.codeFocusIndicator.setTop(cell.layoutInfo.chatHeight);
 			this.codeFocusIndicator.setHeight(cell.layoutInfo.codeIndicatorHeight);
-			this.outputFocusIndicator.setTop(cell.layoutInfo.chatHeight);
 			this.outputFocusIndicator.setHeight(Math.max(cell.layoutInfo.outputIndicatorHeight - cell.viewContext.notebookOptions.getLayoutConfiguration().focusIndicatorGap, 0));
 			this.bottom.domNode.style.transform = `translateY(${cell.layoutInfo.totalHeight - bottomToolbarDimensions.bottomToolbarGap - layoutInfo.cellBottomMargin}px)`;
 		}
@@ -100,8 +97,9 @@ export class CellFocusIndicator extends CellContentPart {
 	}
 
 	private updateFocusIndicatorsForTitleMenu(): void {
-		this.left.domNode.style.transform = `translateY(${this.getIndicatorTopMargin()}px)`;
-		this.right.domNode.style.transform = `translateY(${this.getIndicatorTopMargin()}px)`;
+		const y = (this.currentCell?.layoutInfo.chatHeight ?? 0) + this.getIndicatorTopMargin();
+		this.left.domNode.style.transform = `translateY(${y}px)`;
+		this.right.domNode.style.transform = `translateY(${y}px)`;
 	}
 
 	private getIndicatorTopMargin() {

src/vs/workbench/contrib/notebook/browser/view/cellParts/chat/cellChatController.ts

@@ -8,7 +8,7 @@ import { CancelablePromise, Queue, createCancelablePromise, raceCancellationErro
 import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
 import { Event } from 'vs/base/common/event';
 import { MarkdownString } from 'vs/base/common/htmlContent';
-import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { MovingAverage } from 'vs/base/common/numbers';
 import { StopWatch } from 'vs/base/common/stopwatch';
 import { assertType } from 'vs/base/common/types';
@@ -65,6 +65,7 @@ export class NotebookCellChatController extends Disposable {
 	private _toolbar: MenuWorkbenchToolBar | undefined;
 	private readonly _ctxVisible: IContextKey<boolean>;
 	private readonly _ctxCellWidgetFocused: IContextKey<boolean>;
+	private _widgetDisposableStore: DisposableStore = this._register(new DisposableStore());
 	constructor(
 		private readonly _notebookEditor: INotebookEditorDelegate,
 		private readonly _chatPart: ICellChatPart,
@@ -97,6 +98,10 @@ export class NotebookCellChatController extends Disposable {
 					feedbackMenuId: MENU_CELL_CHAT_WIDGET_FEEDBACK
 				});
 
+				this._widgetDisposableStore.add(this._widget.onDidChangeHeight(() => {
+					this._updateHeight();
+				}));
+
 				this._partContainer.appendChild(this._widget.domNode);
 				this._partContainer.appendChild(this._toolbarDOM.editorToolbar);
 
@@ -123,11 +128,15 @@ export class NotebookCellChatController extends Disposable {
 			this._sessionCtor = undefined;
 		}
 
-		if (this._widget) {
-			this._partContainer.removeChild(this._widget.domNode);
-		}
+		try {
+			if (this._widget) {
+				this._partContainer.removeChild(this._widget.domNode);
+			}
 
-		this._partContainer.removeChild(this._toolbarDOM.editorToolbar);
+			this._partContainer.removeChild(this._toolbarDOM.editorToolbar);
+		} catch (_ex) {
+			// might not be attached
+		}
 
 		this._inlineChatListener?.dispose();
 		this._toolbar?.dispose();
@@ -141,12 +150,25 @@ export class NotebookCellChatController extends Disposable {
 	layout() {
 		if (this._isVisible && this._widget) {
 			const innerEditorWidth = this._cell.layoutInfo.editorWidth;
-			const height = 82 + 8 * 2 /* vertical margin*/;
-
+			const height = this._widget.getHeight();
 			this._widget.layout(new Dimension(innerEditorWidth, height));
 		}
 	}
 
+	private _updateHeight() {
+		const margin = 6;
+		const heightWithMargin = this._isVisible && this._widget
+			? (this._widget.getHeight() - 8 /** shadow */ - 18 /** padding */ + margin /** bottom margin for the cell part */)
+			: 0;
+
+		if (this._cell.chatHeight === heightWithMargin) {
+			return;
+		}
+
+		this._cell.chatHeight = heightWithMargin;
+		this._partContainer.style.height = `${heightWithMargin - margin}px`;
+	}
+
 	async show() {
 		this._isVisible = true;
 		this._partContainer.style.display = 'flex';
@@ -154,7 +176,7 @@ export class NotebookCellChatController extends Disposable {
 		this._widget?.updateInfo(localize('welcome.1', "AI-generated code may be incorrect"));
 		this._ctxVisible.set(true);
 		this._ctxCellWidgetFocused.set(true);
-		this._cell.chatHeight = 82 + 8 * 2 /* vertical margin*/;
+		this._updateHeight();
 
 		this._sessionCtor = createCancelablePromise<void>(async token => {
 			if (this._cell.editorAttached) {
@@ -336,7 +358,7 @@ export class NotebookCellChatController extends Disposable {
 		this._ctxCellWidgetFocused.set(false);
 		this._ctxVisible.set(false);
 		this._widget?.reset();
-		this._cell.chatHeight = 0;
+		this._updateHeight();
 	}
 
 	private async _makeChanges(editor: IActiveCodeEditor, edits: TextEdit[], opts: ProgressingEditsOptions | undefined) {