diff --git a/src/vs/workbench/contrib/notebook/browser/controller/chat/cellChatActions.ts b/src/vs/workbench/contrib/notebook/browser/controller/chat/cellChatActions.ts
index 430062261c9fc..faca219b69cf4 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/chat/cellChatActions.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/chat/cellChatActions.ts
@@ -209,7 +209,7 @@ registerAction2(class extends NotebookAction {
 	}
 
 	async runWithContext(accessor: ServicesAccessor, context: INotebookActionContext) {
-		NotebookChatController.get(context.notebookEditor)?.dismiss();
+		NotebookChatController.get(context.notebookEditor)?.dismiss(false);
 	}
 });
 
diff --git a/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts b/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts
index 93a98b32d60e1..d4b59189404a4 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts
@@ -702,7 +702,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 			this._inlineChatSessionService.releaseSession(this._activeSession);
 		} catch (_err) { }
 
-		this.dismiss();
+		this.dismiss(false);
 	}
 
 	async focusAbove() {
@@ -772,7 +772,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		this._strategy?.cancel();
 		this._activeRequestCts?.cancel();
 		this._widget?.discardChange();
-		this.dismiss();
+		this.dismiss(true);
 	}
 
 	async feedbackLast(kind: InlineChatResponseFeedbackKind) {
@@ -783,25 +783,25 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	}
 
 
-	dismiss() {
-		// move focus back to the cell above
-		if (this._widget) {
-			const widgetIndex = this._widget.afterModelPosition;
-			const currentFocus = this._notebookEditor.getFocus();
-
-			if (currentFocus.start === widgetIndex && currentFocus.end === widgetIndex) {
-				// focus is on the widget
-				if (widgetIndex === 0) {
-					// on top of all cells
-					if (this._notebookEditor.getLength() > 0) {
-						this._notebookEditor.focusNotebookCell(this._notebookEditor.cellAt(0)!, 'container');
-					}
-				} else {
-					const cell = this._notebookEditor.cellAt(widgetIndex - 1);
-					if (cell) {
-						this._notebookEditor.focusNotebookCell(cell, 'container');
-					}
-				}
+	dismiss(discard: boolean) {
+		const widget = this._widget;
+		const widgetIndex = widget?.afterModelPosition;
+		const currentFocus = this._notebookEditor.getFocus();
+		const isWidgetFocused = currentFocus.start === widgetIndex && currentFocus.end === widgetIndex;
+
+		if (widget && isWidgetFocused) {
+			// change focus only when the widget is focused
+			const editingCell = widget.getEditingCell();
+			const shouldFocusEditingCell = editingCell && !discard;
+			const shouldFocusTopCell = widgetIndex === 0 && this._notebookEditor.getLength() > 0;
+			const shouldFocusAboveCell = widgetIndex !== 0 && this._notebookEditor.cellAt(widgetIndex - 1);
+
+			if (shouldFocusEditingCell) {
+				this._notebookEditor.focusNotebookCell(editingCell, 'container');
+			} else if (shouldFocusTopCell) {
+				this._notebookEditor.focusNotebookCell(this._notebookEditor.cellAt(0)!, 'container');
+			} else if (shouldFocusAboveCell) {
+				this._notebookEditor.focusNotebookCell(this._notebookEditor.cellAt(widgetIndex - 1)!, 'container');
 			}
 		}
 
@@ -815,8 +815,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	}
 
 	public override dispose(): void {
-		this.dismiss();
-
+		this.dismiss(false);
 		super.dispose();
 	}
 }
diff --git a/src/vs/workbench/contrib/notebook/browser/view/notebookCellListView.ts b/src/vs/workbench/contrib/notebook/browser/view/notebookCellListView.ts
index e7f7d018a80fc..399934d7c496a 100644
--- a/src/vs/workbench/contrib/notebook/browser/view/notebookCellListView.ts
+++ b/src/vs/workbench/contrib/notebook/browser/view/notebookCellListView.ts
@@ -295,8 +295,20 @@ export class NotebookCellListView<T> extends ListView<T> {
 	}
 
 	removeWhitespace(id: string): void {
-		this.notebookRangeMap.removeWhitespace(id);
-		this.eventuallyUpdateScrollDimensions();
+		const scrollTop = this.scrollTop;
+		const previousRenderRange = this.getRenderRange(this.lastRenderTop, this.lastRenderHeight);
+		const currentPosition = this.notebookRangeMap.getWhitespacePosition(id);
+
+		if (currentPosition > scrollTop) {
+			this.notebookRangeMap.removeWhitespace(id);
+			this.render(previousRenderRange, scrollTop, this.lastRenderHeight, undefined, undefined, false);
+			this._rerender(scrollTop, this.renderHeight, false);
+			this.eventuallyUpdateScrollDimensions();
+		} else {
+			this.notebookRangeMap.removeWhitespace(id);
+			this.eventuallyUpdateScrollDimensions();
+		}
+
 	}
 
 	getWhitespacePosition(id: string): number {