@marimo-team/islands 0.25.3-dev191 → 0.25.3-dev193

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

@@ -21,7 +21,6 @@ import {
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { store } from "@/core/state/jotai";
 import { activeDocumentModelAtom } from "@/core/session/active-notebook";
-import { MultiColumn } from "@/utils/id-tree";
 import { isRtcEnabled } from "../../rtc/state";
 import {
   createNotebookActions,
@@ -206,7 +205,7 @@ describe("applyDocument", () => {
     expect(next.cellHandles[b]).toBeUndefined();
   });
 
-  it("regroups columns from the cells' configs, keeping column ids and collapsed sections", () => {
+  it("regroups columns from the cells' configs, keeping column ids", () => {
     const column1 = { ...cell(d).config, column: 1 };
     const state = notebookOf([
       cell(a),
@@ -214,22 +213,13 @@ describe("applyDocument", () => {
       cell(c),
       cell(d, { config: column1 }),
     ]);
-    const [first, second] = state.cellIds.getColumns();
-    const collapsed = {
-      ...state,
-      cellIds: new MultiColumn([first.collapse(a, b), second]),
-    };
-
     const next = applyDocument(
-      collapsed,
+      state,
       documentOf([cell(c), cell(a), cell(b), cell(d, { config: column1 })]),
     );
 
     expect(next.cellIds.inOrderIds).toEqual([c, a, b, d]);
-    expect(next.cellIds.getColumnIds()).toEqual(
-      collapsed.cellIds.getColumnIds(),
-    );
-    expect(next.cellIds.atOrThrow(0).isCollapsed(a)).toBe(true);
+    expect(next.cellIds.getColumnIds()).toEqual(state.cellIds.getColumnIds());
     expect(next.cellIds.atOrThrow(1).inOrderIds).toEqual([d]);
   });
 
@@ -293,25 +283,6 @@ describe("documentEditMiddleware", () => {
     expect(calls).toEqual([["setCode", a, "a = 2"]]);
   });
 
-  it("sets the order when a cell moves", () => {
-    actions.moveCell({ cellId: c, before: true });
-
-    expect(calls).toEqual([["setOrder", [a, c, b]]]);
-  });
-
-  it("sets the column of a cell that changes column", () => {
-    state = notebookOf([
-      cell(a),
-      cell(b),
-      cell(c, { config: { ...cell(c).config, column: 1 } }),
-    ]);
-
-    actions.moveCell({ cellId: b, direction: "right" });
-
-    expect(state.cellIds.atOrThrow(1).inOrderIds).toEqual([b, c]);
-    expect(calls).toEqual([["updateConfig", b, { column: 1 }]]);
-  });
-
   it("clears columns merged into one", () => {
     state = notebookOf([
       cell(a),

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

@@ -0,0 +1,57 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { createElement, type PropsWithChildren } from "react";
+import { describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import {
+  hideShownCodeAtom,
+  isCodeHidden,
+  restoredEditorStatesAtom,
+  showHiddenCodeAtom,
+  shownHiddenCodeAtom,
+  takeRestoredEditorStateAtom,
+  useRestoredEditorState,
+} from "../editor";
+
+const a = cellId("a");
+const b = cellId("b");
+
+describe("hidden code showing for now", () => {
+  it("hides code unless the cell is showing it", () => {
+    const shown = new Set([a]);
+    expect(isCodeHidden(true, shown, a)).toBe(false);
+    expect(isCodeHidden(true, shown, b)).toBe(true);
+    expect(isCodeHidden(false, shown, b)).toBe(false);
+    expect(isCodeHidden(undefined, shown, b)).toBe(false);
+  });
+
+  it("shows a cell's code and hides it again, saying whether it was showing", () => {
+    const store = createStore();
+    store.set(showHiddenCodeAtom, { cellId: a });
+    expect(store.get(shownHiddenCodeAtom).has(a)).toBe(true);
+    expect(store.set(hideShownCodeAtom, { cellId: a })).toBe(true);
+    expect(store.get(shownHiddenCodeAtom).has(a)).toBe(false);
+    expect(store.set(hideShownCodeAtom, { cellId: a })).toBe(false);
+  });
+});
+
+describe("restored editor states", () => {
+  it("hands a state over once", () => {
+    const store = createStore();
+    store.set(restoredEditorStatesAtom, new Map([[a, { doc: "x = 1" }]]));
+    const wrapper = ({ children }: PropsWithChildren) =>
+      createElement(Provider, { store }, children);
+    const { result } = renderHook(() => useRestoredEditorState(a), {
+      wrapper,
+    });
+    expect(result.current).toEqual({ doc: "x = 1" });
+
+    act(() => store.set(takeRestoredEditorStateAtom, { cellId: a }));
+
+    expect(result.current).toBeNull();
+    expect(store.get(restoredEditorStatesAtom).size).toBe(0);
+  });
+});

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

@@ -0,0 +1,82 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A store whose notebook follows a document model, as the session wires it. */
+
+import { history } from "@codemirror/commands";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { createStore } from "jotai";
+import { mockDocumentModel } from "@/__mocks__/document";
+import {
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
+import { exportedForTesting, notebookAtom, notebookReducer } from "../cells";
+import { followDocument } from "../document";
+import { layoutAtom } from "../layout";
+
+export function setupNotebook(
+  stops: Array<() => void>,
+  cells?: Record<string, Protocol.Cell | null>,
+) {
+  const store = createStore();
+  store.set(notebookAtom, exportedForTesting.initialNotebookState());
+  const { model, stop } = mockDocumentModel(
+    store,
+    cells ? withCells(documentOf(), cells) : documentOf(),
+  );
+  const controller = new AbortController();
+  stops.push(stop, () => controller.abort());
+  followDocument({
+    store,
+    model,
+    reconcile: (document) =>
+      store.set(notebookAtom, (state) =>
+        notebookReducer(state, {
+          type: "reconcileDocument",
+          payload: document,
+        }),
+      ),
+    signal: controller.signal,
+  });
+  return {
+    store,
+    model,
+    order: () => store.get(notebookAtom).cellIds.inOrderIds,
+    /** The cells of each column, as the page lays them out. */
+    columns: () => store.get(layoutAtom).columns.map((column) => column.ids),
+    notebook: () => store.get(notebookAtom),
+    cellIn: (id: string) => store.get(model.documentAtom)?.cells.get(id),
+    /** Give a cell an editor, as mounting it would. */
+    mountEditor: (id: string) => {
+      const view = new EditorView({
+        state: EditorState.create({
+          doc: store.get(model.documentAtom)?.cells.get(id)?.code ?? "",
+          extensions: [history()],
+        }),
+      });
+      store.get(notebookAtom).cellHandles[id as never].current = {
+        editorView: view,
+        editorViewOrNull: view,
+      };
+      return view;
+    },
+    /** Give a cell's output a heading, as the kernel's outline would. */
+    heading: (id: string, level: number) =>
+      store.set(notebookAtom, (state) => ({
+        ...state,
+        cellRuntime: {
+          ...state.cellRuntime,
+          [id]: {
+            ...state.cellRuntime[id as never],
+            outline: { items: [{ name: `H${level}`, level, by: { id } }] },
+          },
+        },
+      })),
+    visible: () => store.get(layoutAtom).columns[0].visible,
+    /** Apply a reducer action, as `useReducerActions()` does. */
+    dispatch: (action: Parameters<typeof notebookReducer>[1]) =>
+      store.set(notebookAtom, (state) => notebookReducer(state, action)),
+  };
+}

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

@@ -0,0 +1,189 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createStore } from "jotai";
+import { describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type { EditableCell } from "@/core/notebook/document";
+import type { CellColumnId } from "@/utils/id-tree";
+import { exportedForTesting, notebookAtom, notebookReducer } from "../cells";
+import { applyDocument } from "../document";
+import type { CellId } from "../ids";
+import {
+  collapsedAtom,
+  columnOf,
+  deriveLayout,
+  foldSections,
+  foldsAfter,
+  foldsHiding,
+  layoutAtom,
+  visibleNeighbor,
+} from "../layout";
+import type { Outline } from "../outline";
+
+const [h1, a, h2, b, c, h1b, d] = ["h1", "a", "h2", "b", "c", "h1b", "d"].map(
+  cellId,
+);
+const ids = [h1, a, h2, b, c, h1b, d];
+
+function heading(level: number): Outline {
+  return { items: [{ name: `H${level}`, level, by: { id: `h${level}` } }] };
+}
+const outlines = new Map<CellId, Outline>([
+  [h1, heading(1)],
+  [h2, heading(2)],
+  [h1b, heading(1)],
+]);
+const outlineOf = (id: CellId) => outlines.get(id) ?? null;
+const column = (id: string, cells: CellId[]) => ({
+  id: id as CellColumnId,
+  ids: cells,
+});
+
+describe("foldSections", () => {
+  it("folds a header until the next header of its level or higher", () => {
+    const folds = foldSections(ids, ids.map(outlineOf), new Set([h2]));
+    expect([...folds]).toEqual([[h2, [b, c]]]);
+  });
+
+  it("hides a folded section inside a folded section", () => {
+    const folds = foldSections(ids, ids.map(outlineOf), new Set([h1, h2]));
+    expect(folds.get(h1)).toEqual([a, h2, b, c]);
+    expect(folds.get(h2)).toEqual([b, c]);
+  });
+
+  it("folds the last section to the end", () => {
+    const folds = foldSections(ids, ids.map(outlineOf), new Set([h1b]));
+    expect(folds.get(h1b)).toEqual([d]);
+  });
+
+  it("ignores cells without a heading", () => {
+    const folds = foldSections([h1, a], [heading(1), null], new Set([a]));
+    expect(folds.size).toBe(0);
+  });
+
+  it("keeps an empty fold for a header with nothing under it yet", () => {
+    const folds = foldSections(
+      [h1, a, h1b],
+      [heading(1), null, heading(1)],
+      new Set([h1b]),
+    );
+    expect(folds.get(h1b)).toEqual([]);
+  });
+});
+
+describe("deriveLayout", () => {
+  const layout = deriveLayout(
+    [column("left", ids), column("right", [])],
+    outlineOf,
+    new Set([h2]),
+  );
+
+  it("lists visible cells per column and keeps empty columns", () => {
+    expect(layout.columns.map((c) => c.id)).toEqual(["left", "right"]);
+    expect(layout.columns[0].visible).toEqual([h1, a, h2, h1b, d]);
+    expect(layout.columns[0].ids).toEqual(ids);
+    expect(layout.columns[1].visible).toEqual([]);
+  });
+
+  it("finds a cell's column, the folds hiding it, and its visible neighbours", () => {
+    expect(columnOf(layout, b).id).toBe("left");
+    expect(() => columnOf(layout, cellId("missing"))).toThrow(
+      "not in the layout",
+    );
+    expect(foldsHiding(layout, c)).toEqual([h2]);
+    expect(foldsHiding(layout, a)).toEqual([]);
+    expect(visibleNeighbor(layout, h2, 1)).toBe(h1b);
+    expect(visibleNeighbor(layout, h1b, -1)).toBe(h2);
+    expect(visibleNeighbor(layout, h1, -1)).toBeUndefined();
+    expect(visibleNeighbor(layout, b, 1)).toBeUndefined();
+  });
+
+  it("names the folds a cell placed after another would fall into", () => {
+    // After the folded header or the last cell of its fold, inside it.
+    expect(foldsAfter(layout, h2)).toEqual([h2]);
+    expect(foldsAfter(layout, c)).toEqual([h2]);
+    // After a visible cell no fold reaches, nothing.
+    expect(foldsAfter(layout, a)).toEqual([]);
+    expect(foldsAfter(layout, h1b)).toEqual([]);
+
+    const nested = deriveLayout(
+      [column("left", ids)],
+      outlineOf,
+      new Set([h1, h2]),
+    );
+    expect(foldsAfter(nested, c)).toEqual([h1, h2]);
+  });
+});
+
+describe("layoutAtom", () => {
+  /** A store whose notebook holds `ids`, in that order, with nothing run. */
+  function notebookOf(...ids: CellId[]) {
+    const store = createStore();
+    const cells = ids.map((id): EditableCell => ({
+      id,
+      code: "",
+      name: "_",
+      config: {
+        column: null,
+        disabled: false,
+        hide_code: false,
+        expand_output: false,
+      },
+    }));
+    store.set(
+      notebookAtom,
+      applyDocument(exportedForTesting.initialNotebookState(), {
+        order: ids,
+        cells: new Map(cells.map((cell) => [cell.id, cell])),
+      }),
+    );
+    return store;
+  }
+
+  /** Deliver `html` as a cell's output, as the kernel's message does. */
+  function output(
+    store: ReturnType<typeof createStore>,
+    id: CellId,
+    html: string,
+  ) {
+    store.set(notebookAtom, (state) =>
+      notebookReducer(state, {
+        type: "handleCellMessage",
+        payload: {
+          cell_id: id,
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: html,
+            timestamp: 1,
+          },
+          console: null,
+          status: "idle",
+          stale_inputs: null,
+          timestamp: 1,
+        },
+      }),
+    );
+  }
+
+  it("is the same layout after output that changes no heading", () => {
+    const store = notebookOf(h1, a);
+    output(store, h1, "<h1>Title</h1>");
+    const layout = store.get(layoutAtom);
+
+    output(store, a, "<p>1</p>");
+    output(store, h1, "<h1>Title</h1>");
+
+    expect(store.get(layoutAtom)).toBe(layout);
+  });
+
+  it("follows a heading arriving in a cell's output", () => {
+    const store = notebookOf(h1, a);
+    store.set(collapsedAtom, new Set([h1]));
+    expect(store.get(layoutAtom).columns[0].folds.size).toBe(0);
+
+    output(store, h1, "<h1>Title</h1>");
+
+    expect(store.get(layoutAtom).columns[0].folds.get(h1)).toEqual([a]);
+  });
+});

package/src/core/cells/__tests__/notebook-outline.test.ts

@@ -80,7 +80,13 @@ describe("notebookOutline", () => {
 
   it("keeps headings in cell order after moving and deleting cells", () => {
     const { store, actions } = setup();
-    actions.moveCell({ cellId: second, before: true });
+    actions.reconcileDocument({
+      order: [second, first],
+      cells: new Map([
+        [first, { id: first, code: "x = 1", name: "_", config: CONFIG }],
+        [second, { id: second, code: "y = 2", name: "_", config: CONFIG }],
+      ]),
+    });
     expect(store.get(notebookOutline)).toEqual({
       items: [...secondOutline.items, ...firstOutline.items],
     });