@marimo-team/islands 0.25.3-dev17 → 0.25.3-dev173

package/src/core/cells/__tests__/cells.test.ts

@@ -6,8 +6,8 @@ import { EditorView } from "@codemirror/view";
 import { createStore } from "jotai";
 import {
   afterAll,
-  afterEach,
   beforeAll,
+  afterEach,
   beforeEach,
   describe,
   expect,
@@ -34,7 +34,6 @@ import {
   type NotebookState,
   notebookAtom,
 } from "../cells";
-import { exportedForTesting as documentTransactionTestExports } from "../document-changes";
 import {
   focusAndScrollCellIntoView,
   scrollToBottom,
@@ -156,10 +155,6 @@ describe("cell reducer", () => {
     firstCellId = state.cellIds.inOrderIds[0];
   });
 
-  afterEach(() => {
-    documentTransactionTestExports.cancelPendingChanges();
-  });
-
   afterAll(() => {
     CellId.create = originalCreate;
   });
@@ -1579,7 +1574,9 @@ describe("cell reducer", () => {
       response: secret,
     });
 
-    expect(cells[0].consoleOutputs).toEqual([{ ...prompt, response: "" }]);
+    expect(cells[0].consoleOutputs).toEqual([
+      { ...prompt, id: 1, response: "" },
+    ]);
     expect(JSON.stringify(cells[0].consoleOutputs)).not.toContain(secret);
   });
 
@@ -1759,187 +1756,6 @@ describe("cell reducer", () => {
     expect(state.cellData[firstCellId].name).toBe("Test Cell");
   });
 
-  it("can set cell IDs and codes", () => {
-    const newIds = [cellId("3"), cellId("4"), cellId("5")];
-    const newCodes = ["code1", "code2", "code3"];
-
-    actions.setCellIds({ cellIds: newIds });
-    expect(state.cellIds.atOrThrow(FIRST_COLUMN).topLevelIds).toEqual(newIds);
-
-    // When codeIsStale is false, lastCodeRun should match code
-    actions.setCellCodes({
-      codes: newCodes,
-      ids: newIds,
-      codeIsStale: false,
-    });
-    newIds.forEach((id, index) => {
-      expect(state.cellData[id].code).toBe(newCodes[index]);
-      expect(state.cellData[id].lastCodeRun).toBe(newCodes[index]);
-      expect(state.cellData[id].edited).toBe(false);
-    });
-
-    // When codeIsStale is true, lastCodeRun should not change
-    const staleCodes = ["stale1", "stale2", "stale3"];
-    actions.setCellCodes({
-      codes: staleCodes,
-      ids: newIds,
-      codeIsStale: true,
-    });
-    newIds.forEach((id, index) => {
-      expect(state.cellData[id].code).toBe(staleCodes[index]);
-      expect(state.cellData[id].lastCodeRun).toBe(newCodes[index]);
-      expect(state.cellData[id].edited).toBe(true);
-    });
-  });
-
-  it("can can add a new cell with/without stale code", () => {
-    actions.setCellCodes({
-      codes: ["new code"],
-      ids: [cellId("2")],
-      codeIsStale: false,
-    });
-
-    expect(state.cellData[cellId("2")].code).toBe("new code");
-    expect(state.cellData[cellId("2")].edited).toBe(false);
-    expect(state.cellData[cellId("2")].lastCodeRun).toBe("new code");
-
-    actions.setCellCodes({
-      codes: ["new code 2"],
-      ids: [cellId("9")],
-      codeIsStale: true,
-    });
-
-    expect(state.cellData[cellId("9")].code).toBe("new code 2");
-    expect(state.cellData[cellId("9")].edited).toBe(true);
-    expect(state.cellData[cellId("9")].lastCodeRun).toBe(null);
-  });
-
-  it("can partial update cell codes", () => {
-    actions.createNewCell({ cellId: firstCellId, before: false });
-    actions.createNewCell({ cellId: cellId("1"), before: false });
-
-    expect(state.cellIds.inOrderIds).toEqual([
-      cellId("0"),
-      cellId("1"),
-      cellId("2"),
-    ]);
-    expect(state.cellData[cellId("0")].code).toBe("");
-    expect(state.cellData[cellId("1")].code).toBe("");
-    expect(state.cellData[cellId("2")].code).toBe("");
-
-    // Update cell 1
-    actions.setCellCodes({
-      codes: ["new code 2"],
-      ids: [cellId("1")],
-      codeIsStale: false,
-    });
-
-    expect(state.cellIds.inOrderIds).toEqual([
-      cellId("0"),
-      cellId("1"),
-      cellId("2"),
-    ]);
-    expect(state.cellData[cellId("0")].code).toBe("");
-    expect(state.cellData[cellId("1")].code).toBe("new code 2");
-    expect(state.cellData[cellId("1")].edited).toBe(false);
-    expect(state.cellData[cellId("2")].code).toBe("");
-  });
-
-  it("can set cell codes with new cell ids, while preserving the old cell data", () => {
-    actions.setCellCodes({
-      codes: ["code1", "code2", "code3"],
-      ids: [cellId("3"), cellId("4"), cellId("5")],
-      codeIsStale: false,
-    });
-    expect(state.cellData[cellId("3")].code).toBe("code1");
-    expect(state.cellData[cellId("4")].code).toBe("code2");
-    expect(state.cellData[cellId("5")].code).toBe("code3");
-
-    // Update with some new cell ids and some old cell ids
-    actions.setCellIds({
-      cellIds: [cellId("1"), cellId("2"), cellId("3"), cellId("4")],
-    });
-    actions.setCellCodes({
-      codes: ["new1", "new2", "code1", "code2"],
-      ids: [cellId("1"), cellId("2"), cellId("3"), cellId("4")],
-      codeIsStale: false,
-    });
-    expect(state.cellData[cellId("1")].code).toBe("new1");
-    expect(state.cellData[cellId("2")].code).toBe("new2");
-    expect(state.cellData[cellId("3")].code).toBe("code1");
-    expect(state.cellData[cellId("4")].code).toBe("code2");
-    expect(state.cellIds.inOrderIds).toEqual([
-      cellId("1"),
-      cellId("2"),
-      cellId("3"),
-      cellId("4"),
-    ]);
-    // Cell 5 data is preserved (possibly used for tracing), but it's not in the cellIds
-    expect(state.cellData[cellId("5")]).not.toBeUndefined();
-  });
-
-  it("can set cell codes with names and configs", () => {
-    const newIds = [cellId("3"), cellId("4")];
-    actions.setCellIds({ cellIds: newIds });
-    actions.setCellCodes({
-      codes: ["code1", "code2"],
-      ids: newIds,
-      codeIsStale: false,
-      names: ["setup_cell", "analysis"],
-      configs: [
-        { hide_code: true, disabled: false, column: null },
-        { hide_code: false, disabled: true, column: null },
-      ],
-    });
-
-    expect(state.cellData[cellId("3")].name).toBe("setup_cell");
-    expect(state.cellData[cellId("3")].config.hide_code).toBe(true);
-    expect(state.cellData[cellId("3")].config.disabled).toBe(false);
-
-    expect(state.cellData[cellId("4")].name).toBe("analysis");
-    expect(state.cellData[cellId("4")].config.hide_code).toBe(false);
-    expect(state.cellData[cellId("4")].config.disabled).toBe(true);
-  });
-
-  it("can set cell codes without names/configs (backward compat)", () => {
-    const newIds = [cellId("3")];
-    actions.setCellIds({ cellIds: newIds });
-    actions.setCellCodes({
-      codes: ["code1"],
-      ids: newIds,
-      codeIsStale: false,
-    });
-
-    // Should use defaults when names/configs not provided
-    expect(state.cellData[cellId("3")].code).toBe("code1");
-    expect(state.cellData[cellId("3")].config.hide_code).toBe(false);
-    expect(state.cellData[cellId("3")].config.disabled).toBe(false);
-  });
-
-  it("can update names and configs on existing cells via setCellCodes", () => {
-    // Set initial state
-    actions.setCellCodes({
-      codes: ["x = 1"],
-      ids: [firstCellId],
-      codeIsStale: false,
-      names: ["old_name"],
-      configs: [{ hide_code: false, disabled: false, column: null }],
-    });
-    expect(state.cellData[firstCellId].name).toBe("old_name");
-    expect(state.cellData[firstCellId].config.hide_code).toBe(false);
-
-    // Update with new name and config (same code)
-    actions.setCellCodes({
-      codes: ["x = 1"],
-      ids: [firstCellId],
-      codeIsStale: true,
-      names: ["new_name"],
-      configs: [{ hide_code: true, disabled: false, column: null }],
-    });
-    expect(state.cellData[firstCellId].name).toBe("new_name");
-    expect(state.cellData[firstCellId].config.hide_code).toBe(true);
-  });
-
   it("can fold and unfold all cells", () => {
     actions.foldAll();
     expect(foldAllBulk).toHaveBeenCalled();
@@ -2630,116 +2446,6 @@ describe("cell reducer", () => {
     expect(state.cellIds.getColumns()[2].topLevelIds).toEqual([cellId("2")]);
   });
 
-  it("rebuildCellColumns regroups cells by config.column", () => {
-    // Create four cells in a single column.
-    actions.createNewCell({ cellId: firstCellId, before: false });
-    actions.createNewCell({ cellId: cellId("1"), before: false });
-    actions.createNewCell({ cellId: cellId("2"), before: false });
-    expect(state.cellIds.getColumns().length).toBe(1);
-
-    // Explicitly set config.column on each cell. This ONLY updates metadata —
-    // updateCellConfig does not touch the MultiColumn tree.
-    actions.updateCellConfig({
-      cellId: cellId("0"),
-      config: { column: 0 },
-    });
-    actions.updateCellConfig({
-      cellId: cellId("2"),
-      config: { column: 1 },
-    });
-
-    // Tree is still a single column at this point.
-    expect(state.cellIds.getColumns().length).toBe(1);
-
-    // Now rebuild the column tree from metadata.
-    actions.rebuildCellColumns({
-      cellIds: [cellId("0"), cellId("1"), cellId("2"), cellId("3")],
-    });
-
-    expect(state.cellIds.getColumns().length).toBe(2);
-    // Cell 1 inherits from 0 (col 0), cell 3 inherits from 2 (col 1).
-    expect(state.cellIds.getColumns()[0].topLevelIds).toEqual([
-      cellId("0"),
-      cellId("1"),
-    ]);
-    expect(state.cellIds.getColumns()[1].topLevelIds).toEqual([
-      cellId("2"),
-      cellId("3"),
-    ]);
-  });
-
-  it("rebuildCellColumns with explicit column on every cell", () => {
-    actions.createNewCell({ cellId: firstCellId, before: false });
-    actions.createNewCell({ cellId: cellId("1"), before: false });
-    actions.createNewCell({ cellId: cellId("2"), before: false });
-
-    for (const [id, col] of [
-      ["0", 1],
-      ["1", 0],
-      ["2", 1],
-      ["3", 0],
-    ] as const) {
-      actions.updateCellConfig({ cellId: cellId(id), config: { column: col } });
-    }
-
-    actions.rebuildCellColumns({
-      cellIds: [cellId("0"), cellId("1"), cellId("2"), cellId("3")],
-    });
-
-    expect(state.cellIds.getColumns().length).toBe(2);
-    expect(state.cellIds.getColumns()[0].topLevelIds).toEqual([
-      cellId("1"),
-      cellId("3"),
-    ]);
-    expect(state.cellIds.getColumns()[1].topLevelIds).toEqual([
-      cellId("0"),
-      cellId("2"),
-    ]);
-  });
-
-  it("rebuildCellColumns collapses to one column when all cells are col 0", () => {
-    actions.createNewCell({ cellId: firstCellId, before: false });
-    actions.addColumnBreakpoint({ cellId: cellId("1") });
-    expect(state.cellIds.getColumns().length).toBe(2);
-
-    // Wipe column metadata back to 0 for both cells.
-    actions.updateCellConfig({ cellId: cellId("0"), config: { column: 0 } });
-    actions.updateCellConfig({ cellId: cellId("1"), config: { column: 0 } });
-
-    actions.rebuildCellColumns({
-      cellIds: [cellId("0"), cellId("1")],
-    });
-
-    expect(state.cellIds.getColumns().length).toBe(1);
-    expect(state.cellIds.getColumns()[0].topLevelIds).toEqual([
-      cellId("0"),
-      cellId("1"),
-    ]);
-  });
-
-  it("rebuildCellColumns follows the provided order, not current tree order", () => {
-    // Start with a single column.
-    actions.createNewCell({ cellId: firstCellId, before: false });
-    actions.createNewCell({ cellId: cellId("1"), before: false });
-    actions.updateCellConfig({ cellId: cellId("0"), config: { column: 0 } });
-    actions.updateCellConfig({ cellId: cellId("2"), config: { column: 1 } });
-
-    // Pass the reversed order. The rebuild should honor it.
-    actions.rebuildCellColumns({
-      cellIds: [cellId("2"), cellId("1"), cellId("0")],
-    });
-
-    // cellIds iteration:
-    //   2 → col=1, pushed to col1, prev=1
-    //   1 → col=null, pushed to col[prev=1] = col1
-    //   0 → col=0, pushed to col0
-    expect(state.cellIds.getColumns()[0].topLevelIds).toEqual([cellId("0")]);
-    expect(state.cellIds.getColumns()[1].topLevelIds).toEqual([
-      cellId("2"),
-      cellId("1"),
-    ]);
-  });
-
   it("can clear output of a single cell", () => {
     // Set up initial state with output
     actions.handleCellMessage({
@@ -3886,31 +3592,9 @@ describe("createTracebackInfoAtom", () => {
   });
 });
 
-describe("setCells snapshot preservation", () => {
+describe("kernel-ready", () => {
   const CELL_A = cellId("A");
   const CELL_B = cellId("B");
-  const newCells: CellData[] = [
-    {
-      id: CELL_A,
-      name: "a",
-      code: "1",
-      edited: false,
-      lastCodeRun: null,
-      lastExecutionTime: null,
-      config: { hide_code: false, disabled: false, column: null },
-      serializedEditorState: null,
-    },
-    {
-      id: CELL_B,
-      name: "b",
-      code: "2",
-      edited: false,
-      lastCodeRun: null,
-      lastExecutionTime: null,
-      config: { hide_code: false, disabled: false, column: null },
-      serializedEditorState: null,
-    },
-  ];
 
   const hydratedState = () =>
     MockNotebook.notebookState({
@@ -3940,44 +3624,97 @@ describe("setCells snapshot preservation", () => {
       },
     });
 
-  beforeEach(async () => {
-    const { isWasm } = await import("@/core/wasm/utils");
-    vi.mocked(isWasm).mockReturnValue(true);
+  it("setLastRun records what the kernel last ran", () => {
+    const state = hydratedState();
+    const next = exportedForTesting.reducer(state, {
+      type: "setLastRun",
+      payload: {
+        lastCodeRun: { [CELL_A]: "0" },
+        lastExecutionTime: { [CELL_A]: 5 },
+      },
+    });
+
+    expect(next.cellData[CELL_A]).toMatchObject({
+      lastCodeRun: "0",
+      lastExecutionTime: 5,
+      edited: true,
+    });
+    expect(next.cellData[CELL_B]).toMatchObject({
+      lastCodeRun: null,
+      lastExecutionTime: null,
+      edited: false,
+    });
+    expect(next.cellRuntime).toBe(state.cellRuntime);
   });
 
-  it("preserves hydrated output in WASM", () => {
+  it("resetCellRuntimes starts every cell's runtime over", () => {
     const next = exportedForTesting.reducer(hydratedState(), {
-      type: "setCells",
-      payload: newCells,
+      type: "resetCellRuntimes",
+      payload: undefined,
     });
 
-    expect(next.cellRuntime[CELL_A].output).toMatchObject({
-      data: "hydrated-A",
+    expect(next.cellRuntime[CELL_A].output).toBeNull();
+    expect(next.cellRuntime[CELL_B].consoleOutputs).toEqual([]);
+  });
+
+  it("resetCellRuntimes starts the scratch cell over too", () => {
+    const state = MockNotebook.notebookState({
+      cellData: { [CELL_A]: { id: CELL_A, code: "1" } },
+      cellRuntime: {
+        [SCRATCH_CELL_ID]: {
+          output: {
+            channel: "output",
+            mimetype: "text/plain",
+            data: "scratch",
+            timestamp: 0,
+          },
+        },
+      },
+    });
+    const next = exportedForTesting.reducer(state, {
+      type: "resetCellRuntimes",
+      payload: undefined,
     });
+
+    expect(next.cellRuntime[SCRATCH_CELL_ID].output).toBeNull();
   });
 
-  it("preserves console-only hydration in WASM", () => {
-    const next = exportedForTesting.reducer(hydratedState(), {
-      type: "setCells",
-      payload: newCells,
+  describe("resetCellRuntimes in WASM", () => {
+    beforeEach(async () => {
+      const { isWasm } = await import("@/core/wasm/utils");
+      vi.mocked(isWasm).mockReturnValue(true);
     });
 
-    expect(next.cellRuntime[CELL_B].consoleOutputs).toHaveLength(1);
-    expect(next.cellRuntime[CELL_B].consoleOutputs[0]).toMatchObject({
-      data: "hydrated-B-stdout",
+    afterEach(async () => {
+      const { isWasm } = await import("@/core/wasm/utils");
+      vi.mocked(isWasm).mockReturnValue(false);
     });
-  });
 
-  it("resets cells with no prior runtime even in WASM", () => {
-    const empty = MockNotebook.notebookState({ cellData: {} });
-    const next = exportedForTesting.reducer(empty, {
-      type: "setCells",
-      payload: newCells,
+    it("keeps hydrated outputs and console", () => {
+      const next = exportedForTesting.reducer(hydratedState(), {
+        type: "resetCellRuntimes",
+        payload: undefined,
+      });
+
+      expect(next.cellRuntime[CELL_A].output).toMatchObject({
+        data: "hydrated-A",
+      });
+      expect(next.cellRuntime[CELL_B].consoleOutputs[0]).toMatchObject({
+        data: "hydrated-B-stdout",
+      });
     });
 
-    expect(next.cellRuntime[CELL_A].output).toBeNull();
-    expect(next.cellRuntime[CELL_A].consoleOutputs).toEqual([]);
-    expect(next.cellRuntime[CELL_B].output).toBeNull();
-    expect(next.cellRuntime[CELL_B].consoleOutputs).toEqual([]);
+    it("starts cells with no prior runtime over", () => {
+      const state = MockNotebook.notebookState({
+        cellData: { [CELL_A]: { id: CELL_A, code: "1" } },
+      });
+      const next = exportedForTesting.reducer(state, {
+        type: "resetCellRuntimes",
+        payload: undefined,
+      });
+
+      expect(next.cellRuntime[CELL_A].output).toBeNull();
+      expect(next.cellRuntime[CELL_A].consoleOutputs).toEqual([]);
+    });
   });
 });

package/src/core/cells/__tests__/collapseConsoleOutputs.test.ts

@@ -2,6 +2,7 @@
 
 import { describe, expect, it } from "vitest";
 import { StatefulOutputMessage } from "@/components/editor/output/ansi-reduce";
+import type { ConsoleOutputMessage } from "@/core/cells/types";
 import type { OutputMessage } from "@/core/kernel/messages";
 import {
   collapseConsoleOutputs,
@@ -279,6 +280,59 @@ describe("collapseConsoleOutputs", () => {
     expect(result[0].data).toContain("Streaming output truncated");
   });
 
+  it("keeps the first id when collapsing outputs", () => {
+    const consoleOutputs = [
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "Hello ",
+        timestamp: 0,
+        id: 7,
+      },
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "World",
+        timestamp: 0,
+        id: 8,
+      },
+    ] satisfies ConsoleOutputMessage[];
+    const result = collapseConsoleOutputs(consoleOutputs);
+    expect(result).toHaveLength(1);
+    expect(result[0].data).toBe("Hello World");
+    expect((result[0] as ConsoleOutputMessage).id).toBe(7);
+  });
+
+  it("keeps ids on outputs that survive truncation", () => {
+    const consoleOutputs = [
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "one\ntwo",
+        timestamp: 0,
+        id: 1,
+      },
+      {
+        mimetype: "text/plain",
+        channel: "stderr",
+        data: "three\nfour",
+        timestamp: 0,
+        id: 2,
+      },
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "five",
+        timestamp: 0,
+        id: 3,
+      },
+    ] satisfies ConsoleOutputMessage[];
+    const result = collapseConsoleOutputs(consoleOutputs, 3);
+    const ids = result.map((output) => (output as ConsoleOutputMessage).id);
+    expect(result[0].data).toContain("Streaming output truncated");
+    expect(ids).toEqual([undefined, 2, 3]);
+  });
+
   describe("ANSI escape sequences", () => {
     it("should handle cursor movement with collapse", () => {
       const consoleOutputs: OutputMessage[] = [