@marimo-team/islands 0.25.3-dev197 → 0.25.3-dev198

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

@@ -1,12 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { createStore } from "jotai";
-import { describe, expect, it } from "vitest";
+import { afterEach, 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 { CellColumnId } from "@/core/cells/layout";
 import type { CellId } from "../ids";
 import {
   collapsedAtom,
@@ -15,10 +11,13 @@ import {
   foldSections,
   foldsAfter,
   foldsHiding,
+  isColumnId,
   layoutAtom,
+  newColumnId,
   visibleNeighbor,
 } from "../layout";
 import type { Outline } from "../outline";
+import { setupNotebook } from "./harness";
 
 const [h1, a, h2, b, c, h1b, d] = ["h1", "a", "h2", "b", "c", "h1b", "d"].map(
   cellId,
@@ -71,6 +70,13 @@ describe("foldSections", () => {
   });
 });
 
+describe("column ids", () => {
+  it("are told apart from cell ids, as drag and drop needs", () => {
+    expect(isColumnId(newColumnId())).toBe(true);
+    expect(isColumnId(cellId("Hbol"))).toBe(false);
+  });
+});
+
 describe("deriveLayout", () => {
   const layout = deriveLayout(
     [column("left", ids), column("right", [])],
@@ -116,74 +122,56 @@ describe("deriveLayout", () => {
 });
 
 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;
-  }
+  const stops: Array<() => void> = [];
+  afterEach(() => {
+    for (const stop of stops.splice(0)) {
+      stop();
+    }
+  });
+  const [c1, c2] = ["c1", "c2"].map(cellId);
 
   /** Deliver `html` as a cell's output, as the kernel's message does. */
   function output(
-    store: ReturnType<typeof createStore>,
+    dispatch: ReturnType<typeof setupNotebook>["dispatch"],
     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,
+    dispatch({
+      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 { store, dispatch } = setupNotebook(stops);
+    output(dispatch, c1, "<h1>Title</h1>");
     const layout = store.get(layoutAtom);
 
-    output(store, a, "<p>1</p>");
-    output(store, h1, "<h1>Title</h1>");
+    output(dispatch, c2, "<p>1</p>");
+    output(dispatch, c1, "<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]));
+    const { store, dispatch } = setupNotebook(stops);
+    store.set(collapsedAtom, new Set([c1]));
     expect(store.get(layoutAtom).columns[0].folds.size).toBe(0);
 
-    output(store, h1, "<h1>Title</h1>");
+    output(dispatch, c1, "<h1>Title</h1>");
 
-    expect(store.get(layoutAtom).columns[0].folds.get(h1)).toEqual([a]);
+    expect(store.get(layoutAtom).columns[0].folds.get(c1)).toEqual([c2]);
   });
 });

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

@@ -6,7 +6,6 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
 import { cellId } from "@/__tests__/branded";
 import { parseOutline } from "@/core/dom/outline";
-import { visibleForTesting } from "@/utils/id-tree";
 import { invariant } from "@/utils/invariant";
 import { Logger } from "@/utils/Logger";
 import { SETUP_CELL_ID } from "../ids";
@@ -35,7 +34,6 @@ function cellIn(result: ReturnType<typeof readExport>, id: string) {
 describe("readExport", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    visibleForTesting.reset();
   });
 
   // Test data factories

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

@@ -1,15 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* oxlint-disable typescript/no-explicit-any */
 import { describe, expect, it } from "vitest";
-import {
-  getCellConfigs,
-  initialNotebookState,
-  type NotebookState,
-} from "@/core/cells/cells";
+import { initialNotebookState, type NotebookState } from "@/core/cells/cells";
 import { CellId } from "@/core/cells/ids";
-import type { CellData, CellRuntimeState } from "@/core/cells/types";
+import type { CellData } from "@/core/cells/types";
 import type { OutputMessage } from "@/core/kernel/messages";
-import { MultiColumn } from "@/utils/id-tree";
 import {
   disabledCellIds,
   enabledCellIds,
@@ -18,103 +13,10 @@ import {
   staleCellIds,
 } from "../utils";
 
-describe("getCellConfigs", () => {
-  it("should return correct cell configs with column indices", () => {
-    // Create mock cell IDs
-    const cellId1 = CellId.create();
-    const cellId2 = CellId.create();
-    const cellId3 = CellId.create();
-    const cellId4 = CellId.create();
-
-    // Create a mock NotebookState
-    const mockState: NotebookState = {
-      ...initialNotebookState(),
-      cellIds: MultiColumn.from([
-        [cellId1, cellId2],
-        [cellId3, cellId4],
-      ]),
-      cellData: {
-        [cellId1]: {
-          id: cellId1,
-          config: { hide_code: false, disabled: false },
-        } as CellData,
-        [cellId2]: {
-          id: cellId2,
-          config: { hide_code: true, disabled: false },
-        } as CellData,
-        [cellId3]: {
-          id: cellId3,
-          config: { hide_code: false, disabled: true },
-        } as CellData,
-        [cellId4]: {
-          id: cellId4,
-          config: { hide_code: true, disabled: true },
-        } as CellData,
-      },
-      cellRuntime: {} as Record<CellId, CellRuntimeState>,
-    };
-
-    // Call the function
-    const result = getCellConfigs(mockState);
-
-    // Assert the results
-    expect(result).toEqual([
-      { hide_code: false, disabled: false, column: 0 },
-      { hide_code: true, disabled: false, column: null },
-      { hide_code: false, disabled: true, column: 1 },
-      { hide_code: true, disabled: true, column: null },
-    ]);
-
-    // Check that the original state was not modified
-    expect(mockState.cellData[cellId1].config).toEqual({
-      hide_code: false,
-      disabled: false,
-    });
-    expect(mockState.cellData[cellId2].config).toEqual({
-      hide_code: true,
-      disabled: false,
-    });
-    expect(mockState.cellData[cellId3].config).toEqual({
-      hide_code: false,
-      disabled: true,
-    });
-    expect(mockState.cellData[cellId4].config).toEqual({
-      hide_code: true,
-      disabled: true,
-    });
-  });
-
-  it("should handle single column", () => {
-    const cellId1 = CellId.create();
-    const mockState: NotebookState = {
-      ...initialNotebookState(),
-      cellIds: MultiColumn.from([[cellId1]]),
-      cellData: {
-        [cellId1]: {
-          id: cellId1,
-          config: { hide_code: false, disabled: false },
-        } as CellData,
-      },
-    };
-
-    const result = getCellConfigs(mockState);
-    expect(result).toEqual([
-      { hide_code: false, disabled: false, column: null },
-    ]);
-  });
-
-  it("should handle empty notebook state", () => {
-    const mockState: NotebookState = initialNotebookState();
-    const result = getCellConfigs(mockState);
-
-    expect(result).toEqual([]);
-  });
-});
-
 describe("staleCellIds", () => {
   it("should return cell IDs that have not been instantiated when auto-instantiate is false", () => {
     const state: NotebookState = {
-      cellIds: MultiColumn.from([["cell1", "cell2"]]),
+      cellIds: ["cell1", "cell2"],
       cellData: {
         cell1: {
           lastExecutionTime: null,
@@ -154,7 +56,7 @@ describe("staleCellIds", () => {
 
   it("should use lastExecutionTime when runElapsedTimeMs is null", () => {
     const state: NotebookState = {
-      cellIds: MultiColumn.from([["cell1", "cell2"]]),
+      cellIds: ["cell1", "cell2"],
       cellData: {
         cell1: {
           lastExecutionTime: null,
@@ -194,7 +96,7 @@ describe("staleCellIds", () => {
 
   it("should return cell IDs that are edited", () => {
     const state: NotebookState = {
-      cellIds: MultiColumn.from([["cell1", "cell2"]]),
+      cellIds: ["cell1", "cell2"],
       cellData: {
         cell1: {
           lastExecutionTime: 123,
@@ -234,7 +136,7 @@ describe("staleCellIds", () => {
 
   it("should return cell IDs that are interrupted", () => {
     const state: NotebookState = {
-      cellIds: MultiColumn.from([["cell1", "cell2"]]),
+      cellIds: ["cell1", "cell2"],
       cellData: {
         cell1: {
           lastExecutionTime: 123,
@@ -274,7 +176,7 @@ describe("staleCellIds", () => {
 
   it("should not return cell IDs that are disabled", () => {
     const state: NotebookState = {
-      cellIds: MultiColumn.from([["cell1", "cell2"]]),
+      cellIds: ["cell1", "cell2"],
       cellData: {
         cell1: {
           lastExecutionTime: 123,
@@ -394,7 +296,7 @@ describe("disabledCellIds", () => {
 
     const state: NotebookState = {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[cellId1, cellId2, cellId3]]),
+      cellIds: [cellId1, cellId2, cellId3],
       cellData: {
         [cellId1]: { id: cellId1, config: { disabled: true } } as CellData,
         [cellId2]: { id: cellId2, config: { disabled: false } } as CellData,
@@ -415,7 +317,7 @@ describe("enabledCellIds", () => {
 
     const state: NotebookState = {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[cellId1, cellId2, cellId3]]),
+      cellIds: [cellId1, cellId2, cellId3],
       cellData: {
         [cellId1]: { id: cellId1, config: { disabled: true } } as CellData,
         [cellId2]: { id: cellId2, config: { disabled: false } } as CellData,

package/src/core/cells/actions.ts

@@ -3,21 +3,28 @@
 /** Notebook operations that edit the document and then adjust the editor. */
 
 import { historyField } from "@codemirror/commands";
-import { atom, type Getter, useSetAtom } from "jotai";
+import { atom, type Getter, type Setter, useSetAtom } from "jotai";
 import { useMemo } from "react";
 import { scrollAndHighlightCell } from "@/components/editor/links/cell-link";
 import type { EditableDocument } from "@/core/notebook/document";
-import type { CellColumnId } from "@/utils/id-tree";
 import { Logger } from "@/utils/Logger";
 import { clamp } from "@/utils/math";
+import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
+import { splitEditor } from "../codemirror/language/utils";
 import type { SerializedEditorState } from "../codemirror/types";
 import { canCollapseOutline } from "../dom/outline";
 import type { CellConfig } from "../network/types";
 import { store } from "../state/jotai";
 import {
+  addColumnAtom,
+  addColumnBreakpointAtom,
+  compactColumnsAtom,
+  deleteColumnAtom,
   dropCellOverCellAtom,
   dropCellOverColumnAtom,
+  mergeAllColumnsAtom,
   moveCellAtom,
+  moveColumnAtom,
   moveCellsRelativeToAtom,
   moveCellToIndexAtom,
   placeCells,
@@ -27,6 +34,7 @@ import {
 import {
   type DeletedCell,
   notebookAtom,
+  notebookReducer,
   releaseCellAtoms,
   useReducerActions,
 } from "./cells";
@@ -39,8 +47,10 @@ import {
 } from "./editor";
 import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
 import {
+  type CellColumnId,
   collapsedAtom,
   columnOf,
+  findColumn,
   foldOf,
   foldsAfter,
   foldsHiding,
@@ -91,7 +101,7 @@ export const createCellAtom = atom(
     if (options.skipIfCodeExists && hasCellWithCode(document, code)) {
       return null;
     }
-    const { cellIds } = get(notebookAtom);
+    const layout = get(layoutAtom);
 
     // The document knows cells and order. Columns come from each cell's
     // `column` config, which the cell before it provides unless it starts a
@@ -100,31 +110,31 @@ export const createCellAtom = atom(
     let column: number | null = null;
     // The cell the new one comes right after, whose folds would hide it.
     let follows: CellId | undefined;
-    const layout = get(layoutAtom);
     if (options.after) {
       // A folded header hides the cells after it, so a cell placed after
       // the header goes after its fold.
       follows = foldOf(layout, options.after)?.at(-1) ?? options.after;
       placement = { after: follows };
     } else if (options.before) {
-      const tree = cellIds.findWithId(options.before);
-      const index = cellIds.indexOf(tree);
+      const target = columnOf(layout, options.before);
+      const index = layout.columns.indexOf(target);
       placement = { before: options.before };
-      follows = tree.before(options.before);
-      if (index > 0 && tree.indexOfOrThrow(options.before) === 0) {
+      follows = target.ids[target.ids.indexOf(options.before) - 1];
+      if (index > 0 && target.ids[0] === options.before) {
         column = index;
       }
     } else {
-      const tree =
-        (options.endOf && cellIds.get(options.endOf)) ?? cellIds.atOrThrow(0);
-      if (tree.length > 0) {
-        follows = tree.last();
-        placement = { after: follows };
-      } else if (cellIds.colLength > 1) {
-        column = cellIds.indexOf(tree);
+      const target =
+        (options.endOf && findColumn(layout, options.endOf)) ??
+        layout.columns[0];
+      const last = target.ids.at(-1);
+      if (last !== undefined) {
+        follows = last;
+        placement = { after: last };
+      } else if (layout.columns.length > 1) {
+        column = layout.columns.indexOf(target);
       }
     }
-
     const id = options.id ?? CellId.create();
     const config = { ...options.config, column };
     set(modelOf(get).editAtom, (tx) => {
@@ -157,6 +167,70 @@ export const createCellAtom = atom(
   },
 );
 
+/**
+ * Split a cell at its cursor. The text before stays, the text after
+ * becomes a new cell below, which is focused. Returns the new cell's id,
+ * or null for a cell without an editor.
+ */
+export const splitCellAtom = atom(
+  null,
+  (get, set, { cellId }: { cellId: CellId }): CellId | null => {
+    const view = get(notebookAtom).cellHandles[cellId]?.current?.editorView;
+    if (!view) {
+      return null;
+    }
+    const { beforeCursorCode, afterCursorCode } = splitEditor(view);
+    const update = prepareCellCodeUpdate(view.state, beforeCursorCode);
+    if (update) {
+      view.dispatch(update);
+    }
+    const id = CellId.create();
+    set(modelOf(get).editAtom, (tx) => {
+      tx.setCode(cellId, beforeCursorCode);
+      tx.createCell({ id, code: afterCursorCode }, { after: cellId });
+    });
+    clearOutput(set, cellId);
+    focusWhenMounted(set, id);
+    return id;
+  },
+);
+
+/** Undo a split: the cell takes its code back and the cell split off goes. */
+export const undoSplitCellAtom = atom(
+  null,
+  (
+    get,
+    set,
+    {
+      cellId,
+      newCellId,
+      snapshot,
+    }: { cellId: CellId; newCellId: CellId; snapshot: string },
+  ): void => {
+    const view = get(notebookAtom).cellHandles[cellId]?.current?.editorView;
+    const update = view && prepareCellCodeUpdate(view.state, snapshot);
+    if (view && update) {
+      view.dispatch(update);
+    }
+    const document = documentOf(get);
+    set(modelOf(get).editAtom, (tx) => {
+      tx.setCode(cellId, snapshot);
+      if (document.cells.has(newCellId)) {
+        tx.deleteCell(newCellId);
+      }
+    });
+    releaseCellAtoms(newCellId);
+    clearOutput(set, cellId);
+  },
+);
+
+/** The output belonged to the code before the change; neither version has run. */
+function clearOutput(set: Setter, cellId: CellId): void {
+  set(notebookAtom, (state) =>
+    notebookReducer(state, { type: "clearCellOutput", payload: { cellId } }),
+  );
+}
+
 /**
  * Focus the cell before or after `cellId`. Past the edge of the column,
  * create one there, unless `create` is false.
@@ -283,9 +357,7 @@ export const collapseAllCellsAtom = atom(null, (get, set): void => {
   set(
     collapsedAtom,
     new Set(
-      cellIds.inOrderIds.filter((id) =>
-        canCollapseOutline(cellRuntime[id]?.outline),
-      ),
+      cellIds.filter((id) => canCollapseOutline(cellRuntime[id]?.outline)),
     ),
   );
 });
@@ -318,18 +390,18 @@ export const addSetupCellAtom = atom(
   null,
   (get, set, { code }: { code?: string } = {}): void => {
     const { cellIds } = get(notebookAtom);
-    if (cellIds.setupCellExists()) {
+    if (cellIds.includes(SETUP_CELL_ID)) {
       focusWhenMounted(set, SETUP_CELL_ID);
       return;
     }
     // The setup cell leads the first column, which may be empty.
-    const first = cellIds.atOrThrow(0);
+    const first = get(layoutAtom).columns[0];
     set(createCellAtom, {
       id: SETUP_CELL_ID,
       name: SETUP_CELL_ID,
       code: code ?? "",
-      ...(first.length > 0
-        ? { before: first.atOrThrow(0) }
+      ...(first.ids.length > 0
+        ? { before: first.ids[0] }
         : { endOf: first.id }),
     });
   },
@@ -528,9 +600,25 @@ export function useCellActions() {
   const dropCellOverColumn = useSetAtom(dropCellOverColumnAtom);
   const sendToTop = useSetAtom(sendToTopAtom);
   const sendToBottom = useSetAtom(sendToBottomAtom);
+  const addColumn = useSetAtom(addColumnAtom);
+  const addColumnBreakpoint = useSetAtom(addColumnBreakpointAtom);
+  const deleteColumn = useSetAtom(deleteColumnAtom);
+  const mergeAllColumns = useSetAtom(mergeAllColumnsAtom);
+  const compactColumns = useSetAtom(compactColumnsAtom);
+  const moveColumn = useSetAtom(moveColumnAtom);
+  const splitCell = useSetAtom(splitCellAtom);
+  const undoSplitCell = useSetAtom(undoSplitCellAtom);
   return useMemo(
     () => ({
       ...operations,
+      splitCell,
+      undoSplitCell,
+      addColumn,
+      addColumnBreakpoint,
+      deleteColumn,
+      mergeAllColumns,
+      compactColumns,
+      moveColumn,
       moveCell,
       moveCellToIndex,
       moveCellsRelativeTo,
@@ -563,9 +651,7 @@ export type CellActions = ReturnType<typeof useCellActions>;
 /** Scroll to the first cell that is currently running. */
 export function notebookScrollToRunning() {
   const { cellIds, cellRuntime } = store.get(notebookAtom);
-  const cellId = cellIds.inOrderIds.find(
-    (id) => cellRuntime[id]?.status === "running",
-  );
+  const cellId = cellIds.find((id) => cellRuntime[id]?.status === "running");
   if (!cellId) {
     return;
   }

package/src/core/cells/add-missing-import.ts

@@ -36,7 +36,7 @@ export function maybeAddMissingImport({
     `import[ \t]+${moduleName}[ \t]+as[ \t]+${variableName}`,
     "g",
   );
-  for (const cell of cellIds.inOrderIds) {
+  for (const cell of cellIds) {
     if (regex.test(cellData[cell].code)) {
       return false;
     }