@marimo-team/islands 0.25.3-dev20 → 0.25.3-dev201

package/src/core/cells/columns.ts

@@ -0,0 +1,38 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** How a document's cells fall into columns. */
+
+import type { EditableDocument } from "@/core/notebook/document";
+import type { CellId } from "./ids";
+
+/**
+ * Each cell's column: its `column` config, or the column of the cell before
+ * it. The first cell starts in column 0.
+ */
+export function resolveColumns(
+  document: EditableDocument,
+): Map<CellId, number> {
+  const columns = new Map<CellId, number>();
+  let current = 0;
+  for (const id of document.order) {
+    const column = document.cells.get(id)?.config.column;
+    if (column != null && column >= 0) {
+      current = column;
+    }
+    // The protocol's ids are plain strings; the notebook brands them.
+    columns.set(id as CellId, current);
+  }
+  return columns;
+}
+
+/** The cells of each column in order, with empty columns for gaps. */
+export function groupByColumn(document: EditableDocument): CellId[][] {
+  const columns: CellId[][] = [[]];
+  for (const [id, column] of resolveColumns(document)) {
+    while (columns.length <= column) {
+      columns.push([]);
+    }
+    columns[column].push(id);
+  }
+  return columns;
+}

package/src/core/cells/document.ts

@@ -0,0 +1,289 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * Keep the notebook state and the host document in step while cells live
+ * in both.
+ *
+ * TODO(MO-8028): Delete this module once components read cells from the
+ * document model and `notebookAtom` holds only runtime and editor state.
+ */
+
+import { dequal as isEqual } from "dequal";
+import type { createStore } from "jotai";
+import { createRef, type RefObject } from "react";
+import type { CellHandle } from "@/components/editor/notebook-cell";
+import type {
+  DocumentEdit,
+  DocumentEditor,
+  DocumentModel,
+  EditableCell,
+  EditableDocument,
+} from "@/core/notebook/document";
+import type { DispatchedActionOf } from "@/utils/createReducer";
+import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
+import { canUseRtc, isRtcEnabled } from "../rtc/state";
+import { store } from "../state/jotai";
+import { activeDocumentModelAtom } from "../session/active-notebook";
+import type { NotebookState, ReducerCellActions } from "./cells";
+import type { CellId } from "./ids";
+import { followColumnsAtom } from "./layout";
+import {
+  type CellData,
+  createCell,
+  createCellConfig,
+  createCellRuntimeState,
+  isEdited,
+} from "./types";
+
+type CellAction = DispatchedActionOf<ReducerCellActions>;
+
+// Following the document ----------------------------------------------------
+
+/** The cell's data with the document's fields, and its editor brought along. */
+function followCell(
+  current: CellData,
+  cell: EditableCell,
+  handle: RefObject<CellHandle | null> | undefined,
+): CellData {
+  // Withheld source keeps the code this client has.
+  const code = cell.code ?? current.code;
+  const config = isEqual(current.config, cell.config)
+    ? current.config
+    : createCellConfig(cell.config);
+  if (
+    code === current.code &&
+    cell.name === current.name &&
+    config === current.config
+  ) {
+    return current;
+  }
+  if (code === current.code) {
+    return { ...current, name: cell.name, config };
+  }
+  // RTC cells follow Loro. Other editors take the host's code; one that
+  // shows it already, as the editor this client types in does, is left alone.
+  if (!isRtcEnabled() || !canUseRtc(current.id)) {
+    const view = handle?.current?.editorViewOrNull;
+    const update = view && prepareCellCodeUpdate(view.state, code);
+    if (update) {
+      view.dispatch(update);
+    }
+  }
+  return {
+    ...current,
+    code,
+    name: cell.name,
+    config,
+    edited: isEdited(current.lastCodeRun, code),
+  };
+}
+
+/**
+ * Make the notebook's cells and order follow the document.
+ *
+ * Returns the same state when nothing differs, so a local edit that the
+ * document already reflects changes nothing.
+ *
+ * TODO(MO-8028): Delete with `reconcileDocument` once cells are read from
+ * the document model.
+ */
+export function applyDocument(
+  state: NotebookState,
+  document: EditableDocument,
+): NotebookState {
+  // 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) {
+    const cell = document.cells.get(id);
+    const current = state.cellData[id];
+    if (cell === undefined) {
+      continue;
+    }
+    if (current === undefined) {
+      // A new cell with code is stale until it runs. The kernel's last run
+      // settles it on kernel-ready.
+      const code = cell.code ?? "";
+      cellData = {
+        ...cellData,
+        [id]: createCell({
+          id,
+          code,
+          name: cell.name,
+          config: createCellConfig(cell.config),
+          edited: code.trim().length > 0,
+        }),
+      };
+      cellRuntime = {
+        ...cellRuntime,
+        [id]: cellRuntime[id] ?? createCellRuntimeState(),
+      };
+      cellHandles = { ...cellHandles, [id]: createRef() };
+      continue;
+    }
+    const next = followCell(current, cell, state.cellHandles[id]);
+    if (next !== current) {
+      cellData = { ...cellData, [id]: next };
+    }
+  }
+  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> =>
+      Object.fromEntries(
+        Object.entries(record).filter(([id]) => !gone.has(id as CellId)),
+      ) as Record<CellId, V>;
+    cellData = keep(cellData);
+    cellRuntime = keep(cellRuntime);
+    cellHandles = keep(cellHandles);
+  }
+  if (
+    cellIds === state.cellIds &&
+    cellData === state.cellData &&
+    removed.length === 0
+  ) {
+    return state;
+  }
+  return { ...state, cellIds, cellData, cellRuntime, cellHandles };
+}
+
+/**
+ * Reconcile the model's document and name its columns, now and on each
+ * change, until the signal aborts.
+ *
+ * TODO(MO-8028): Delete with `applyDocument`.
+ */
+export function followDocument({
+  store,
+  model,
+  reconcile,
+  signal,
+}: {
+  store: ReturnType<typeof createStore>;
+  model: Pick<DocumentModel, "documentAtom">;
+  reconcile: (document: EditableDocument) => void;
+  signal: AbortSignal;
+}): void {
+  const follow = () => {
+    const document = store.get(model.documentAtom);
+    if (document !== null) {
+      reconcile(document);
+      store.set(followColumnsAtom, document);
+    }
+  };
+  const unsubscribe = store.sub(model.documentAtom, follow);
+  signal.addEventListener("abort", unsubscribe, { once: true });
+  follow();
+}
+
+// Editing the document ------------------------------------------------------
+
+/**
+ * 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 diffNotebook(
+  prev: NotebookState,
+  next: NotebookState,
+): DocumentEdit | null {
+  const steps: Array<(tx: DocumentEditor) => void> = [];
+  for (const id of next.cellIds) {
+    const cell = next.cellData[id];
+    const before = prev.cellData[id];
+    if (before !== undefined && cell !== before && before.code !== cell.code) {
+      steps.push((tx) => tx.setCode(id, cell.code));
+    }
+  }
+  if (steps.length === 0) {
+    return null;
+  }
+  return (tx) => {
+    for (const step of steps) {
+      step(tx);
+    }
+  };
+}
+
+/**
+ * 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): Only code is left in `diffNotebook`. The middleware goes
+ * with it, once the editors bind to the model.
+ */
+export function documentEditMiddleware(
+  prev: NotebookState,
+  next: NotebookState,
+  action: CellAction,
+): void {
+  // The document replaces the notebook's cells, so following it is no edit.
+  if (action.type === "reconcileDocument") {
+    return;
+  }
+  if (prev.cellIds === next.cellIds && prev.cellData === next.cellData) {
+    return;
+  }
+  const model = store.get(activeDocumentModelAtom);
+  if (model === null) {
+    return;
+  }
+  const edit = diffNotebook(prev, next);
+  if (edit !== null) {
+    store.set(model.editAtom, edit);
+  }
+}
+
+// Commands that read the document on the host -------------------------------
+
+/**
+ * Send this client's edits to the host before a command that reads them
+ * there. Without a document or a backend, as on a static page, there is
+ * nothing to send and the command proceeds.
+ */
+export async function flushDocument(): Promise<void> {
+  const model = store.get(activeDocumentModelAtom);
+  if (
+    model === null ||
+    store.get(model.syncStatusAtom).status === "unavailable"
+  ) {
+    return;
+  }
+  await model.flush();
+}
+
+/**
+ * Run a legacy save after this client's edits reach the host, and keep
+ * later edits local until it finishes. The save carries the whole
+ * notebook and the host takes it as the document, so an edit sent
+ * meanwhile would be overwritten. When the edits cannot be sent the save
+ * still runs, and the model resumes against the saved document.
+ *
+ * TODO(MO-8028): Goes with the legacy save, along with the hold.
+ */
+export async function withDocumentSave<T>(save: () => Promise<T>): Promise<T> {
+  const model = store.get(activeDocumentModelAtom);
+  if (model === null) {
+    return save();
+  }
+  let flushed = true;
+  try {
+    await model.flush();
+  } catch {
+    flushed = false;
+  }
+  const release = model.hold();
+  let result: T;
+  try {
+    result = await save();
+  } finally {
+    release();
+  }
+  if (!flushed) {
+    void model.retry().catch(() => undefined);
+  }
+  return result;
+}

package/src/core/cells/edit.ts

@@ -0,0 +1,30 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The active document model, as the notebook operations reach it. */
+
+import type { Getter, Setter } from "jotai";
+import type { DocumentModel, EditableDocument } from "@/core/notebook/document";
+import { activeDocumentModelAtom } from "../session/active-notebook";
+import { notebookAtom } from "./cells";
+import type { CellId } from "./ids";
+
+export function modelOf(get: Getter): DocumentModel {
+  const model = get(activeDocumentModelAtom);
+  if (!model) {
+    throw new Error("The notebook has no document to edit");
+  }
+  return model;
+}
+
+export function documentOf(get: Getter): EditableDocument {
+  const document = get(modelOf(get).documentAtom);
+  if (!document) {
+    throw new Error("The notebook has no document to edit");
+  }
+  return document;
+}
+
+/** Scroll to and focus the cell once it mounts. */
+export function focusWhenMounted(set: Setter, cellId: CellId): void {
+  set(notebookAtom, (state) => ({ ...state, scrollKey: cellId }));
+}

package/src/core/cells/editor.ts

@@ -0,0 +1,82 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What only this browser knows about cells: hidden code showing for now, and editor state to restore. */
+
+import { atom, useAtomValue } from "jotai";
+import { useMemo } from "react";
+import type { SerializedEditorState } from "../codemirror/types";
+import type { CellId } from "./ids";
+
+/**
+ * Hidden cells whose code is showing for now: a cell being edited, one just
+ * created, or one just made hidden. Blurring the cell hides it again.
+ */
+export const shownHiddenCodeAtom = atom<ReadonlySet<CellId>>(new Set<CellId>());
+
+export const useShowsHiddenCode = (cellId: CellId) =>
+  useAtomValue(
+    useMemo(
+      () => atom((get) => get(shownHiddenCodeAtom).has(cellId)),
+      [cellId],
+    ),
+  );
+
+/** Whether a cell's code is hidden on the page. */
+export function isCodeHidden(
+  hideCode: boolean | undefined,
+  shown: ReadonlySet<CellId>,
+  cellId: CellId,
+): boolean {
+  return Boolean(hideCode) && !shown.has(cellId);
+}
+
+export const showHiddenCodeAtom = atom(
+  null,
+  (_get, set, { cellId }: { cellId: CellId }): void => {
+    set(shownHiddenCodeAtom, (cells) => new Set([...cells, cellId]));
+  },
+);
+
+/** Hide a cell's code again. Returns whether it was showing. */
+export const hideShownCodeAtom = atom(
+  null,
+  (get, set, { cellId }: { cellId: CellId }): boolean => {
+    if (!get(shownHiddenCodeAtom).has(cellId)) {
+      return false;
+    }
+    set(shownHiddenCodeAtom, (cells) => {
+      const next = new Set(cells);
+      next.delete(cellId);
+      return next;
+    });
+    return true;
+  },
+);
+
+/**
+ * Editor state for cells about to mount again, such as a cell undo brings
+ * back. The editor takes its entry when it mounts, and deleting the cell
+ * drops it.
+ */
+export const restoredEditorStatesAtom = atom<
+  ReadonlyMap<CellId, SerializedEditorState>
+>(new Map<CellId, SerializedEditorState>());
+
+export const useRestoredEditorState = (cellId: CellId) =>
+  useAtomValue(
+    useMemo(
+      () => atom((get) => get(restoredEditorStatesAtom).get(cellId) ?? null),
+      [cellId],
+    ),
+  );
+
+export const takeRestoredEditorStateAtom = atom(
+  null,
+  (_get, set, { cellId }: { cellId: CellId }): void => {
+    set(restoredEditorStatesAtom, (states) => {
+      const next = new Map(states);
+      next.delete(cellId);
+      return next;
+    });
+  },
+);

package/src/core/cells/layout.ts

@@ -0,0 +1,221 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The cells as they appear on the page: by column, with folded sections hidden. */
+
+import { atom, useAtomValue } from "jotai";
+import type { EditableDocument } from "@/core/notebook/document";
+import { findCollapseRange } from "../dom/outline";
+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. */
+  readonly ids: readonly CellId[];
+  /** The cells on the page: those no fold hides. */
+  readonly visible: readonly CellId[];
+  /**
+   * Each folded header and the cells it hides, nested sections included. A
+   * folded header with nothing under it yet has an empty fold.
+   */
+  readonly folds: ReadonlyMap<CellId, readonly CellId[]>;
+}
+
+export interface Layout {
+  readonly columns: readonly LayoutColumn[];
+}
+
+/**
+ * The cells each folded header hides. A section runs from its header to the
+ * cell before the next header of the same or a higher level, so a folded
+ * header inside a folded section is hidden along with its own section.
+ */
+export function foldSections(
+  ids: readonly CellId[],
+  outlines: (Outline | null)[],
+  collapsed: ReadonlySet<CellId>,
+): Map<CellId, readonly CellId[]> {
+  const folds = new Map<CellId, readonly CellId[]>();
+  ids.forEach((id, index) => {
+    const outline = outlines[index];
+    if (!collapsed.has(id) || !outline || outline.items.length === 0) {
+      return;
+    }
+    const range = findCollapseRange(index, outlines);
+    if (range === null) {
+      return;
+    }
+    folds.set(id, ids.slice(index + 1, range[1] + 1));
+  });
+  return folds;
+}
+
+export function deriveLayout(
+  columns: ReadonlyArray<{ id: CellColumnId; ids: readonly CellId[] }>,
+  outlineOf: (id: CellId) => Outline | null,
+  collapsed: ReadonlySet<CellId>,
+): Layout {
+  return {
+    columns: columns.map(({ id, ids }) => {
+      const folds = foldSections(ids, ids.map(outlineOf), collapsed);
+      const hidden = new Set([...folds.values()].flat());
+      return { id, ids, visible: ids.filter((c) => !hidden.has(c)), folds };
+    }),
+  };
+}
+
+/**
+ * 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): 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(
+    columnsOf(document, get(columnIdsAtom), get(emptyColumnIdsAtom)),
+    (id) => outlines[id] ?? null,
+    get(collapsedAtom),
+  );
+});
+
+export const useLayout = () => useAtomValue(layoutAtom);
+
+export const numColumnsAtom = atom((get) => get(layoutAtom).columns.length);
+
+export function findColumn(
+  layout: Layout,
+  columnId: CellColumnId,
+): LayoutColumn | undefined {
+  return layout.columns.find((column) => column.id === columnId);
+}
+
+/** The column holding `id`. Throws for a cell the layout does not have. */
+export function columnOf(layout: Layout, id: CellId): LayoutColumn {
+  const column = layout.columns.find((column) => column.ids.includes(id));
+  if (column === undefined) {
+    throw new Error(`Cell ${id} is not in the layout`);
+  }
+  return column;
+}
+
+/** The cells a folded header hides, or undefined when `id` is not a folded header. */
+export function foldOf(
+  layout: Layout,
+  id: CellId,
+): readonly CellId[] | undefined {
+  for (const column of layout.columns) {
+    const fold = column.folds.get(id);
+    if (fold !== undefined) {
+      return fold;
+    }
+  }
+  return undefined;
+}
+
+/** The headers whose folds hide `id`. */
+export function foldsHiding(layout: Layout, id: CellId): CellId[] {
+  const column = layout.columns.find((column) => column.ids.includes(id));
+  if (column === undefined) {
+    return [];
+  }
+  return [...column.folds]
+    .filter(([, hidden]) => hidden.includes(id))
+    .map(([header]) => header);
+}
+
+/**
+ * The folded headers whose sections a cell without a heading would fall
+ * into, placed right after `id`: those hiding `id`, and `id` itself when it
+ * is a folded header. A section runs until the next heading, so a cell
+ * placed after the last cell of a fold is still inside it.
+ */
+export function foldsAfter(layout: Layout, id: CellId): CellId[] {
+  const hiding = foldsHiding(layout, id);
+  return foldOf(layout, id) === undefined ? hiding : [...hiding, id];
+}
+
+/** The visible cell `offset` places from `id` in its column, if there is one. */
+export function visibleNeighbor(
+  layout: Layout,
+  id: CellId,
+  offset: -1 | 1,
+): CellId | undefined {
+  const { visible } = columnOf(layout, id);
+  const index = visible.indexOf(id);
+  return index === -1 ? undefined : visible[index + offset];
+}