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

package/src/core/cells/document.ts

@@ -8,7 +8,6 @@
  * document model and `notebookAtom` holds only runtime and editor state.
  */
 
-import type * as Protocol from "@marimo-team/client/protocol";
 import { dequal as isEqual } from "dequal";
 import type { createStore } from "jotai";
 import { createRef, type RefObject } from "react";
@@ -21,15 +20,13 @@ import type {
   EditableDocument,
 } from "@/core/notebook/document";
 import type { DispatchedActionOf } from "@/utils/createReducer";
-import { CollapsibleTree, MultiColumn } from "@/utils/id-tree";
 import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
-import type { CellConfig } from "../network/types";
 import { canUseRtc, isRtcEnabled } from "../rtc/state";
 import { store } from "../state/jotai";
 import { activeDocumentModelAtom } from "../session/active-notebook";
 import type { NotebookState, ReducerCellActions } from "./cells";
-import { groupByColumn } from "./columns";
 import type { CellId } from "./ids";
+import { followColumnsAtom } from "./layout";
 import {
   type CellData,
   createCell,
@@ -42,27 +39,6 @@ type CellAction = DispatchedActionOf<ReducerCellActions>;
 
 // Following the document ----------------------------------------------------
 
-/** Rebuild the tree for `columns`, keeping column ids. */
-function reconcileColumns(
-  current: MultiColumn<CellId>,
-  columns: CellId[][],
-): MultiColumn<CellId> {
-  const previous = current.getColumns();
-  // The document cannot hold an empty column, so keep the ones this client has.
-  while (previous[columns.length]?.length === 0) {
-    columns.push([]);
-  }
-  if (
-    columns.length === previous.length &&
-    columns.every((ids, index) => isEqual(ids, previous[index].inOrderIds))
-  ) {
-    return current;
-  }
-  return new MultiColumn(
-    columns.map((ids, index) => CollapsibleTree.from(ids, previous[index]?.id)),
-  );
-}
-
 /** The cell's data with the document's fields, and its editor brought along. */
 function followCell(
   current: CellData,
@@ -115,9 +91,12 @@ export function applyDocument(
   state: NotebookState,
   document: EditableDocument,
 ): NotebookState {
-  const cellIds = reconcileColumns(state.cellIds, groupByColumn(document));
+  // The protocol's ids are plain strings; the notebook brands them.
+  const cellIds = isEqual(document.order, state.cellIds)
+    ? state.cellIds
+    : ([...document.order] as CellId[]);
   let { cellData, cellRuntime, cellHandles } = state;
-  for (const id of cellIds.inOrderIds) {
+  for (const id of cellIds) {
     const cell = document.cells.get(id);
     const current = state.cellData[id];
     if (cell === undefined) {
@@ -149,8 +128,8 @@ export function applyDocument(
       cellData = { ...cellData, [id]: next };
     }
   }
-  const present = new Set(cellIds.inOrderIds);
-  const removed = state.cellIds.inOrderIds.filter((id) => !present.has(id));
+  const present = new Set(cellIds);
+  const removed = state.cellIds.filter((id) => !present.has(id));
   if (removed.length > 0) {
     const gone = new Set(removed);
     const keep = <V>(record: Record<CellId, V>): Record<CellId, V> =>
@@ -172,7 +151,8 @@ export function applyDocument(
 }
 
 /**
- * Reconcile the model's document now and on each change, until the signal aborts.
+ * Reconcile the model's document and name its columns, now and on each
+ * change, until the signal aborts.
  *
  * TODO(MO-8028): Delete with `applyDocument`.
  */
@@ -191,6 +171,7 @@ export function followDocument({
     const document = store.get(model.documentAtom);
     if (document !== null) {
       reconcile(document);
+      store.set(followColumnsAtom, document);
     }
   };
   const unsubscribe = store.sub(model.documentAtom, follow);
@@ -200,107 +181,22 @@ export function followDocument({
 
 // Editing the document ------------------------------------------------------
 
-const CONFIG_FIELDS = ["disabled", "hide_code", "expand_output"] as const;
-
-/** The fields of `after` that differ from `before`, leaving `column` to the tree. */
-function configPatch(
-  before: CellConfig,
-  after: CellConfig,
-): Partial<Protocol.CellConfig> | null {
-  const patch: Partial<Protocol.CellConfig> = {};
-  for (const field of CONFIG_FIELDS) {
-    if (before[field] !== after[field]) {
-      patch[field] = after[field] ?? false;
-    }
-  }
-  return Object.keys(patch).length === 0 ? null : patch;
-}
-
 /**
- * The `column` configs that make the document's columns match the tree's.
- *
- * A cell without a column follows the cell before it, so only cells whose
- * resolved column differs from their position need one.
+ * The edit that writes the code a reducer action changed, or null when none
+ * did. Order, creation, deletion, names and configs are edited on the model
+ * directly by their atoms.
  */
-function columnFixes(state: NotebookState): Array<[CellId, number | null]> {
-  const fixes: Array<[CellId, number | null]> = [];
-  const columns = state.cellIds.getColumns();
-  if (columns.length === 1) {
-    for (const id of columns[0].inOrderIds) {
-      const column = state.cellData[id].config.column;
-      if (column != null && column !== 0) {
-        fixes.push([id, null]);
-      }
-    }
-    return fixes;
-  }
-  let inherited = 0;
-  columns.forEach((column, target) => {
-    for (const id of column.inOrderIds) {
-      const configured = state.cellData[id].config.column;
-      const resolved =
-        configured != null && configured >= 0 ? configured : inherited;
-      if (resolved !== target) {
-        fixes.push([id, target]);
-      }
-      inherited = target;
-    }
-  });
-  return fixes;
-}
-
-/** The edit that makes the document match `next`, or null when it already does. */
 function diffNotebook(
   prev: NotebookState,
   next: NotebookState,
 ): DocumentEdit | null {
   const steps: Array<(tx: DocumentEditor) => void> = [];
-  const prevOrder = prev.cellIds.inOrderIds;
-  const nextOrder = next.cellIds.inOrderIds;
-  const prevIds = new Set(prevOrder);
-  const nextIds = new Set(nextOrder);
-  for (const id of prevOrder) {
-    if (!nextIds.has(id)) {
-      steps.push((tx) => tx.deleteCell(id));
-    }
-  }
-  const created: CellId[] = [];
-  for (const id of nextOrder) {
+  for (const id of next.cellIds) {
     const cell = next.cellData[id];
     const before = prev.cellData[id];
-    if (!prevIds.has(id) || before === undefined) {
-      created.push(id);
-      steps.push((tx) =>
-        tx.createCell({
-          id,
-          code: cell.code,
-          name: cell.name,
-          config: cell.config,
-        }),
-      );
-      continue;
-    }
-    if (before === cell) {
-      continue;
-    }
-    if (before.code !== cell.code) {
+    if (before !== undefined && cell !== before && before.code !== cell.code) {
       steps.push((tx) => tx.setCode(id, cell.code));
     }
-    if (before.name !== cell.name) {
-      steps.push((tx) => tx.setName(id, cell.name));
-    }
-    const patch = configPatch(before.config, cell.config);
-    if (patch !== null) {
-      steps.push((tx) => tx.updateConfig(id, patch));
-    }
-  }
-  for (const [id, column] of columnFixes(next)) {
-    steps.push((tx) => tx.updateConfig(id, { column }));
-  }
-  // Created cells go to the end, so the order is set when they belong elsewhere.
-  const survivors = prevOrder.filter((id) => nextIds.has(id));
-  if (!isEqual(nextOrder, [...survivors, ...created])) {
-    steps.push((tx) => tx.setOrder(nextOrder));
   }
   if (steps.length === 0) {
     return null;
@@ -316,8 +212,8 @@ function diffNotebook(
  * Middleware for the notebook reducer. Writes what an action changed about
  * the cells and their order into the document model as one edit.
  *
- * TODO(MO-8028): Each action family that moves onto the model removes its
- * part of `diffNotebook`. The middleware goes with the last one, code.
+ * TODO(MO-8028): Only code is left in `diffNotebook`. The middleware goes
+ * with it, once the editors bind to the model.
  */
 export function documentEditMiddleware(
   prev: NotebookState,

package/src/core/cells/layout.ts

@@ -3,15 +3,83 @@
 /** The cells as they appear on the page: by column, with folded sections hidden. */
 
 import { atom, useAtomValue } from "jotai";
-import type { CellColumnId } from "@/utils/id-tree";
+import type { EditableDocument } from "@/core/notebook/document";
 import { findCollapseRange } from "../dom/outline";
-import { cellIdsAtom, cellOutlinesAtom } from "./cells";
-import type { CellId } from "./ids";
+import { activeDocumentModelAtom } from "../session/active-notebook";
+import { cellOutlinesAtom } from "./cells";
+import { groupByColumn } from "./columns";
+import { CellId } from "./ids";
 import type { Outline } from "./outline";
 
+/** A column's identity on this page, as drag and drop and saved widths refer to it. */
+export type CellColumnId = string & { __brand: "CellColumnId" };
+
 /** The headers whose sections the user folded. Only this browser knows. */
 export const collapsedAtom = atom<ReadonlySet<CellId>>(new Set<CellId>());
 
+/**
+ * This browser's names for the document's columns, by position. The
+ * document knows columns only by position, so `followColumnsAtom` names
+ * a column as it appears and forgets one as it goes, and the column
+ * operations carry the names along as they rearrange.
+ */
+export const columnIdsAtom = atom<readonly CellColumnId[]>([]);
+
+/**
+ * Trailing columns a rearrangement here emptied. The document cannot hold
+ * an empty column at its end, so the page keeps them, as drop targets for
+ * the cell being dragged, until a drop compacts them.
+ */
+export const emptyColumnIdsAtom = atom<readonly CellColumnId[]>([]);
+
+export function newColumnId(): CellColumnId {
+  return `column-${CellId.create()}` as CellColumnId;
+}
+
+export function isColumnId(id: unknown): id is CellColumnId {
+  return typeof id === "string" && id.startsWith("column-");
+}
+
+/**
+ * Name the document's columns as the document changes. A column that
+ * appears takes over an empty column the page keeps, or gets a fresh id;
+ * one that goes is forgotten.
+ */
+export const followColumnsAtom = atom(
+  null,
+  (get, set, document: EditableDocument): void => {
+    const count = groupByColumn(document).length;
+    const ids = get(columnIdsAtom);
+    if (ids.length === count) {
+      return;
+    }
+    const empty = get(emptyColumnIdsAtom);
+    const named = [...ids, ...empty].slice(0, count);
+    while (named.length < count) {
+      named.push(newColumnId());
+    }
+    set(columnIdsAtom, named);
+    set(emptyColumnIdsAtom, empty.slice(Math.max(0, count - ids.length)));
+  },
+);
+
+/** The document's columns, each with its id, then the empty ones the page keeps. */
+function columnsOf(
+  document: EditableDocument | null,
+  ids: readonly CellColumnId[],
+  empty: readonly CellColumnId[],
+): Array<{ id: CellColumnId; ids: CellId[] }> {
+  const groups = document === null ? [[]] : groupByColumn(document);
+  return [
+    ...groups.map((cells, index) => ({
+      // Before the first document names them, columns go by position.
+      id: ids[index] ?? (`column-${index}` as CellColumnId),
+      ids: cells,
+    })),
+    ...empty.map((id) => ({ id, ids: [] as CellId[] })),
+  ];
+}
+
 export interface LayoutColumn {
   readonly id: CellColumnId;
   /** Every cell of the column, in order. */
@@ -69,18 +137,17 @@ export function deriveLayout(
 }
 
 /**
- * The layout of the notebook's cells.
+ * The layout of the notebook's cells: the document's order and columns,
+ * this browser's column ids and folds, and the kernel's outlines.
  *
- * TODO(MO-8028): Order and columns come from the reducer's tree until the
- * arrange family edits the document; then they derive from the document and
- * the column ids this browser keeps. Outlines come from executions later.
+ * TODO(MO-8028): Outlines come from executions later.
  */
 export const layoutAtom = atom((get): Layout => {
+  const model = get(activeDocumentModelAtom);
+  const document = model === null ? null : get(model.documentAtom);
   const outlines = get(cellOutlinesAtom);
   return deriveLayout(
-    get(cellIdsAtom)
-      .getColumns()
-      .map((column) => ({ id: column.id, ids: column.inOrderIds })),
+    columnsOf(document, get(columnIdsAtom), get(emptyColumnIdsAtom)),
     (id) => outlines[id] ?? null,
     get(collapsedAtom),
   );
@@ -88,6 +155,8 @@ export const layoutAtom = atom((get): Layout => {
 
 export const useLayout = () => useAtomValue(layoutAtom);
 
+export const numColumnsAtom = atom((get) => get(layoutAtom).columns.length);
+
 export function findColumn(
   layout: Layout,
   columnId: CellColumnId,

package/src/core/cells/pending-delete-service.ts

@@ -49,7 +49,7 @@ const { valueAtom: pendingDeleteStateAtom, useActions } = createReducerAndAtoms(
       const variables = store.get(variablesAtom);
 
       const emptyCells = new Set(
-        notebook.cellIds.inOrderIds.filter(
+        notebook.cellIds.filter(
           (id) => notebook.cellData[id].code.trim() === "",
         ),
       );

package/src/core/cells/utils.ts

@@ -25,7 +25,7 @@ export function notebookNeedsRun(state: NotebookState) {
 }
 
 export function notebookCells(state: NotebookState) {
-  return state.cellIds.inOrderIds.map((cellId) => state.cellData[cellId]);
+  return state.cellIds.map((cellId) => state.cellData[cellId]);
 }
 
 export function notebookCellEditorViews({ cellHandles }: NotebookState) {
@@ -42,7 +42,7 @@ export function notebookCellEditorViews({ cellHandles }: NotebookState) {
 export function disabledCellIds(state: NotebookState) {
   const { cellIds, cellData } = state;
   const disabledCells: CellId[] = [];
-  for (const cellId of cellIds.inOrderIds) {
+  for (const cellId of cellIds) {
     const cell = cellData[cellId];
     if (cell.config.disabled) {
       disabledCells.push(cellId);
@@ -54,7 +54,7 @@ export function disabledCellIds(state: NotebookState) {
 export function enabledCellIds(state: NotebookState) {
   const { cellIds, cellData } = state;
   const enabledCells: CellId[] = [];
-  for (const cellId of cellIds.inOrderIds) {
+  for (const cellId of cellIds) {
     const cell = cellData[cellId];
     if (!cell.config.disabled) {
       enabledCells.push(cellId);
@@ -108,7 +108,7 @@ export function getDescendantsStatus(
  */
 export function staleCellIds(state: NotebookState) {
   const { cellIds, cellData, cellRuntime } = state;
-  return cellIds.inOrderIds.filter(
+  return cellIds.filter(
     (cellId) =>
       isUninstantiated({
         // runElapstedTimeMs is what we've seen in this session

package/src/core/codemirror/copilot/__tests__/getCodes.test.ts

@@ -9,7 +9,6 @@ import { createCell } from "@/core/cells/types";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
-import { MultiColumn } from "@/utils/id-tree";
 import { cellConfigExtension } from "../../config/extension";
 import { adaptiveLanguageConfiguration } from "../../language/extension";
 import {
@@ -159,7 +158,7 @@ describe("getCodes", () => {
     ];
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellIds: [Cells.cell1, Cells.cell2],
       cellHandles: {
         [Cells.cell1]: { current: mockEditorViews[0] },
         [Cells.cell2]: { current: mockEditorViews[1] },
@@ -177,7 +176,7 @@ describe("getCodes", () => {
     ];
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellIds: [Cells.cell1, Cells.cell2],
       cellHandles: {
         [Cells.cell1]: { current: mockEditorViews[0] },
         [Cells.cell2]: { current: mockEditorViews[1] },
@@ -196,7 +195,7 @@ describe("getCodes", () => {
     ];
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2, Cells.cell3]]),
+      cellIds: [Cells.cell1, Cells.cell2, Cells.cell3],
       cellHandles: {
         [Cells.cell1]: { current: mockEditorViews[0] },
         [Cells.cell2]: { current: mockEditorViews[1] },
@@ -218,7 +217,7 @@ describe("getCodes", () => {
     const otherCode = "print('Hello World')";
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellIds: [Cells.cell1, Cells.cell2],
       cellData: {
         [Cells.cell1]: createCell({ id: Cells.cell1, code: "import os" }),
         [Cells.cell2]: createCell({ id: Cells.cell2, code: "x = 1" }),
@@ -232,7 +231,7 @@ describe("getCodes", () => {
     const otherCode = "print('Hello World')";
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellIds: [Cells.cell1, Cells.cell2],
       cellData: {
         [Cells.cell1]: createCell({ id: Cells.cell1, code: "import os" }),
         [Cells.cell2]: createCell({ id: Cells.cell2, code: "x = 1" }),
@@ -250,7 +249,7 @@ describe("topologicalCodesAtom", () => {
   it("should include every cell when no editor is mounted", () => {
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2, Cells.cell3]]),
+      cellIds: [Cells.cell1, Cells.cell2, Cells.cell3],
       cellData: {
         [Cells.cell1]: createCell({ id: Cells.cell1, code: "import os" }),
         [Cells.cell2]: createCell({ id: Cells.cell2, code: "x = 1" }),

package/src/core/codemirror/copilot/getCodes.ts

@@ -12,7 +12,7 @@ import { getEditorCodeAsPython } from "../language/utils";
 const notebookCellCodes = atom((get) => {
   const notebook = get(notebookAtom);
   const codes = Objects.fromEntries(
-    notebook.cellIds.inOrderIds.map((id) => {
+    notebook.cellIds.map((id) => {
       // A cell editor mounts progressively, so the notebook store is the only
       // complete source of code until every editor is built.
       const editorView = notebook.cellHandles[id]?.current?.editorViewOrNull;
@@ -54,7 +54,7 @@ export function getOtherCellsCode(otherCode: string) {
 
 const inOrderCellIdsAtom = atom((get) => {
   const notebook = get(notebookAtom);
-  return notebook.cellIds.inOrderIds;
+  return notebook.cellIds;
 });
 const topologicalCellIdsAtom = atom((get) => {
   const cellIds = get(inOrderCellIdsAtom);

package/src/core/codemirror/lsp/__tests__/notebook-lsp.test.ts

@@ -11,7 +11,6 @@ import type { CellId } from "@/core/cells/ids";
 import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
 import { createCell } from "@/core/cells/types";
 import { store } from "@/core/state/jotai";
-import { MultiColumn } from "@/utils/id-tree";
 import { topologicalCodesAtom } from "../../copilot/getCodes";
 import { lspWorkspaceAtom } from "@/core/saving/file-state";
 import { languageAdapterState } from "../../language/extension";
@@ -69,7 +68,7 @@ describe("merged document with unmounted editors", () => {
   it("should include the code of cells whose editors are not mounted", () => {
     store.set(notebookAtom, {
       ...initialNotebookState(),
-      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellIds: [Cells.cell1, Cells.cell2],
       cellData: {
         [Cells.cell1]: createCell({ id: Cells.cell1, code: "import math" }),
         [Cells.cell2]: createCell({

package/src/core/codemirror/lsp/utils.ts

@@ -13,7 +13,7 @@ export function isKnownCellDocumentUri(uri: string): boolean {
   if (!CellDocumentUri.is(uri)) {
     return false;
   }
-  return store.get(cellIdsAtom).inOrderIds.includes(CellDocumentUri.parse(uri));
+  return store.get(cellIdsAtom).includes(CellDocumentUri.parse(uri));
 }
 
 export function getLspRootUri() {

package/src/core/documentation/doc-lookup.ts

@@ -17,7 +17,7 @@ export async function requestOutputDocumentation(
   qualifiedName: string,
 ): Promise<void> {
   // We need any valid cell_id for the completion request.
-  const cellId = getCells().inOrderIds.at(0);
+  const cellId = getCells().at(0);
   if (!cellId) {
     return;
   }

package/src/core/edit-app.tsx

@@ -27,11 +27,8 @@ import { useSetOutputsExpanded } from "../components/editor/actions/useSetOutput
 import { CellArray } from "../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
 import { useHotkey } from "../hooks/useHotkey";
-import {
-  hasCellsAtom,
-  notebookIsRunningAtom,
-  numColumnsAtom,
-} from "./cells/cells";
+import { hasCellsAtom, notebookIsRunningAtom } from "./cells/cells";
+import { numColumnsAtom } from "./cells/layout";
 import { useCellActions } from "./cells/actions";
 import type { AppConfig, UserConfig } from "./config/config-schema";
 import { RuntimeState } from "./kernel/RuntimeState";

package/src/core/errors/error-entries.ts

@@ -80,7 +80,7 @@ export function getCellErrorEntries(store: JotaiStore): CellErrorEntry[] {
   const { cellIds, cellRuntime, cellData } = store.get(notebookAtom);
   const entries: CellErrorEntry[] = [];
 
-  for (const [cellIndex, cellId] of cellIds.inOrderIds.entries()) {
+  for (const [cellIndex, cellId] of cellIds.entries()) {
     const cell = cellRuntime[cellId];
     const output = cell.output;
     if (!output || !isErrorMime(output.mimetype)) {

package/src/core/islands/__tests__/bootstrap.test.ts

@@ -103,15 +103,11 @@ describe("initializeIslands document", () => {
     }
 
     const notebook = store.get(notebookAtom);
-    expect(notebook.cellIds.inOrderIds).toEqual(["c1", "c3", "k7f2"]);
+    expect(notebook.cellIds).toEqual(["c1", "c3", "k7f2"]);
     expect(notebook.cellData[cellId("c3")].code).toBe("x = 21");
     expect(notebook.cellRuntime[cellId("c3")]).toBeDefined();
 
     handlers?.message({ event: "reset", data: {} });
-    expect(store.get(notebookAtom).cellIds.inOrderIds).toEqual([
-      "c1",
-      "c3",
-      "k7f2",
-    ]);
+    expect(store.get(notebookAtom).cellIds).toEqual(["c1", "c3", "k7f2"]);
   });
 });

package/src/core/islands/components/web-components.tsx

@@ -104,7 +104,7 @@ export class MarimoIslandElement extends HTMLElement {
    */
   private getCellIdFromIndex(idx: number): CellId | undefined {
     const { cellIds } = store.get(notebookAtom);
-    return cellIds.inOrderIds.at(idx);
+    return cellIds.at(idx);
   }
 
   /**

package/src/core/notebook/document/__tests__/codemirror.test.ts

@@ -82,6 +82,14 @@ function languages(id: CellId) {
       dropCellOverColumn: vi.fn(),
       sendToTop: vi.fn(),
       sendToBottom: vi.fn(),
+      addColumn: vi.fn(),
+      addColumnBreakpoint: vi.fn(),
+      deleteColumn: vi.fn(),
+      mergeAllColumns: vi.fn(),
+      compactColumns: vi.fn(),
+      moveColumn: vi.fn(),
+      splitCell: vi.fn(),
+      undoSplitCell: vi.fn(),
       onRun: vi.fn(),
       deleteCell: vi.fn(),
       afterToggleMarkdown: vi.fn(),

package/src/core/session/__tests__/useNotebookSession.hook.test.tsx

@@ -272,7 +272,7 @@ describe("useNotebookSession document model", () => {
     });
 
     let notebook = store.get(notebookAtom);
-    expect(notebook.cellIds.inOrderIds).toEqual(["c1", "c3", "k7f2"]);
+    expect(notebook.cellIds).toEqual(["c1", "c3", "k7f2"]);
     expect(notebook.cellData[cellId("k7f2")]).toMatchObject({
       code: "y = x * 2",
       config: { hide_code: true },

package/src/stories/cell.stories.tsx

@@ -15,7 +15,6 @@ import { requestClientAtom } from "@/core/network/requests";
 import { resolveRequestClient } from "@/core/network/resolve";
 import type { CellConfig } from "@/core/network/types";
 import { ConnectionState } from "@/core/session/types";
-import { MultiColumn } from "@/utils/id-tree";
 import type { Milliseconds, Seconds } from "@/utils/time";
 import { Cell as EditorCell } from "../components/editor/notebook-cell";
 import { TooltipProvider } from "../components/ui/tooltip";
@@ -55,7 +54,7 @@ const Cell: React.FC<{
         lastExecutionTime: null,
       },
     },
-    cellIds: MultiColumn.from([[cid]]),
+    cellIds: [cid],
     cellRuntime: {
       [cid]: createCellRuntimeState({
         output: overrides.output ?? null,

package/src/stories/layout/vertical/one-column.stories.tsx

@@ -12,7 +12,6 @@ import { connectionAtom } from "@/core/network/connection";
 import { requestClientAtom } from "@/core/network/requests";
 import { resolveRequestClient } from "@/core/network/resolve";
 import { ConnectionState } from "@/core/session/types";
-import { MultiColumn } from "@/utils/id-tree";
 import type { Milliseconds, Seconds } from "@/utils/time";
 import { CellArray } from "../../../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../../../components/editor/renderers/cells-renderer";
@@ -52,7 +51,7 @@ LongReprA()`,
       lastExecutionTime: null,
     },
   },
-  cellIds: MultiColumn.from([[cellId]]),
+  cellIds: [cellId],
   cellRuntime: {
     [cellId]: createCellRuntimeState({
       output: {

package/src/stories/runtime-story.tsx

@@ -16,7 +16,6 @@ import type { EnvironmentOperation } from "@/core/alerts/environment";
 import { notebookAtom } from "@/core/cells/cells";
 import { CellId } from "@/core/cells/ids";
 import { createCell, createCellRuntimeState } from "@/core/cells/types";
-import { MultiColumn } from "@/utils/id-tree";
 import { userConfigAtom } from "@/core/config/config";
 import {
   type AppConfig,
@@ -111,7 +110,7 @@ function RuntimePreview({
       {props.surface === "notebook" && (
         <div className="flex-1 min-w-0 pt-12">
           <ConnectionNotice appConfig={appConfig} onRetry={reconnect} />
-          {notebook.cellIds.inOrderIds.map((cellId) => (
+          {notebook.cellIds.map((cellId) => (
             <VerticalLayoutWrapper key={cellId} appConfig={appConfig}>
               <Cell
                 cellId={cellId}
@@ -351,7 +350,7 @@ export function RuntimeStory(props: Props) {
       const cellId = CellId.create();
       state.set(notebookAtom, (value) => ({
         ...value,
-        cellIds: MultiColumn.from([[cellId]]),
+        cellIds: [cellId],
         cellData: {
           ...value.cellData,
           [cellId]: createCell({