From 5cd092d0e329ceddaeb8bb38a7d399e80720b4a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Krassowski?=
 <5832902+krassowski@users.noreply.github.com>
Date: Mon, 2 Dec 2024 20:44:21 +0000
Subject: [PATCH] Fix disappearing cells (heal offsets after updating estimated
 sizes) (#17000)

* Add a galata test against #16978

* Fix offsets of items after updating estimated sizes

* Do not touch outputs if cell model is disposed

* Rename `h` to `notebook`
Please enter the commit message for your changes. Lines starting
---
 .../test/jupyterlab/windowed-notebook.test.ts | 49 +++++++++++++++++++
 packages/notebook/src/windowing.ts            |  2 +-
 .../src/components/windowedlist.ts            |  9 +++-
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/galata/test/jupyterlab/windowed-notebook.test.ts b/galata/test/jupyterlab/windowed-notebook.test.ts
index 19835e65aac5..cde2d69e0a97 100644
--- a/galata/test/jupyterlab/windowed-notebook.test.ts
+++ b/galata/test/jupyterlab/windowed-notebook.test.ts
@@ -398,6 +398,55 @@ test('should remove all cells including hidden outputs artifacts', async ({
   expect(found).toEqual(false);
 });
 
+test('should display cells below on scrolling after inserting a cell on top', async ({
+  page,
+  tmpPath
+}) => {
+  // Regression test against "disappearing cells" issue:
+  // https://github.com/jupyterlab/jupyterlab/issues/16978
+  await page.notebook.openByPath(`${tmpPath}/${fileName}`);
+
+  const notebook = await page.notebook.getNotebookInPanelLocator()!;
+  const firstCell = notebook.locator('.jp-Cell[data-windowed-list-index="1"]');
+  const lastCell = notebook.locator('.jp-Cell[data-windowed-list-index="18"]');
+  await firstCell.waitFor();
+
+  const bbox = await notebook.boundingBox();
+  await page.mouse.move(bbox!.x, bbox!.y);
+
+  // Needs to be two separate mouse wheel events.
+  await page.mouse.wheel(0, 3000);
+  await page.mouse.wheel(0, 3000);
+
+  // Scroll down to reveal last cell to ensure these all items have been measured...
+  await Promise.all([
+    firstCell.waitFor({ state: 'hidden' }),
+    lastCell.waitFor()
+  ]);
+
+  await page.mouse.wheel(0, -3000);
+  await page.mouse.wheel(0, -3000);
+
+  // ...then scroll back up and select first cell.
+  await Promise.all([
+    lastCell.waitFor({ state: 'hidden' }),
+    firstCell.waitFor()
+  ]);
+  await page.notebook.selectCells(0);
+
+  // Insert cell below.
+  await page.keyboard.press('b');
+  await page.mouse.wheel(0, 3000);
+  await page.mouse.wheel(0, 3000);
+
+  // Scroll down again.
+  await Promise.all([
+    firstCell.waitFor({ state: 'hidden' }),
+    lastCell.waitFor()
+  ]);
+  await expect(lastCell).toBeVisible();
+});
+
 test('should center on next cell after rendering markdown cell and advancing', async ({
   page,
   tmpPath
diff --git a/packages/notebook/src/windowing.ts b/packages/notebook/src/windowing.ts
index 4b7204f9b88a..8d11d204ecbf 100644
--- a/packages/notebook/src/windowing.ts
+++ b/packages/notebook/src/windowing.ts
@@ -58,7 +58,7 @@ export class NotebookViewModel extends WindowedListModel {
 
     const nLines = model.sharedModel.getSource().split('\n').length;
     let outputsLines = 0;
-    if (model instanceof CodeCellModel) {
+    if (model instanceof CodeCellModel && !model.isDisposed) {
       for (let outputIdx = 0; outputIdx < model.outputs.length; outputIdx++) {
         const output = model.outputs.get(outputIdx);
         const data = output.data['text/plain'];
diff --git a/packages/ui-components/src/components/windowedlist.ts b/packages/ui-components/src/components/windowedlist.ts
index 782cb6ba3f86..bdfdb839c9b5 100644
--- a/packages/ui-components/src/components/windowedlist.ts
+++ b/packages/ui-components/src/components/windowedlist.ts
@@ -663,8 +663,13 @@ export abstract class WindowedListModel implements WindowedList.IModel {
 
         offset += size;
       }
-
-      this._measuredAllUntilIndex = index;
+      // Because the loop above updates estimated sizes,
+      // we need to fix (heal) offsets of the remaining items.
+      for (let i = index + 1; i < this._widgetSizers.length; i++) {
+        const sizer = this._widgetSizers[i];
+        const previous = this._widgetSizers[i - 1];
+        sizer.offset = previous.offset + previous.size;
+      }
     }
 
     for (let i = 0; i <= this._measuredAllUntilIndex; i++) {