@marimo-team/islands 0.25.3-dev30 → 0.25.3-dev302

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,167 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * Make the notebook state follow the host document while cells live in
+ * both. Edits go the other way through the document model directly: the
+ * cell operations write it, and each cell editor is bound to its cell.
+ *
+ * TODO(MO-8349): 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 } from "react";
+import type {
+  DocumentModel,
+  EditableCell,
+  EditableDocument,
+} from "@/core/notebook/document";
+import { store } from "../state/jotai";
+import { activeDocumentModelAtom } from "../session/active-notebook";
+import type { NotebookState } from "./cells";
+import type { CellId } from "./ids";
+import { followColumnsAtom } from "./layout";
+import {
+  type CellData,
+  createCell,
+  createCellConfig,
+  createCellRuntimeState,
+} from "./types";
+
+// Following the document ----------------------------------------------------
+
+/** The cell's data after the document's cell, the same object when nothing differs. */
+function followCell(current: CellData, cell: EditableCell): 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 };
+  }
+  return { ...current, code, name: cell.name, config };
+}
+
+/**
+ * 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-8349): 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) {
+      cellData = {
+        ...cellData,
+        [id]: createCell({
+          id,
+          code: cell.code ?? "",
+          name: cell.name,
+          config: createCellConfig(cell.config),
+        }),
+      };
+      cellRuntime = {
+        ...cellRuntime,
+        [id]: cellRuntime[id] ?? createCellRuntimeState(),
+      };
+      cellHandles = { ...cellHandles, [id]: createRef() };
+      continue;
+    }
+    const next = followCell(current, cell);
+    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-8349): 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();
+}
+
+// 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();
+}

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/focus.ts

@@ -3,7 +3,13 @@
 import type { EditorView } from "@codemirror/view";
 import { atom, useAtomValue } from "jotai";
 import { createReducerAndAtoms } from "@/utils/createReducer";
-import type { CellConfig, RuntimeState } from "../network/types";
+import {
+  type CellExecutionInputs,
+  type CellRunStatus,
+  cellExecutionInputsAtom,
+  cellExecutionView,
+} from "@/core/executions";
+import type { CellConfig } from "../network/types";
 import { type NotebookState, notebookAtom } from "./cells";
 import type { CellId } from "./ids";
 import { SCRATCH_CELL_ID } from "./ids";
@@ -75,7 +81,8 @@ export const lastFocusedCellAtom = atom<{
   name: string;
   config: CellConfig;
   cellId: CellId;
-  status: RuntimeState;
+  status: CellRunStatus;
+  blocked: boolean;
   getEditorView: () => EditorView | null;
   hasOutput: boolean;
   hasConsoleOutput: boolean;
@@ -89,16 +96,28 @@ export const lastFocusedCellAtom = atom<{
     return null;
   }
 
-  return cellFocusDetails(cellId, get(notebookAtom));
+  return cellFocusDetails(
+    cellId,
+    get(notebookAtom),
+    get(cellExecutionInputsAtom),
+  );
 });
 
 export function cellFocusDetailsAtom(cellId: CellId) {
   return atom((get) => {
-    return cellFocusDetails(cellId, get(notebookAtom));
+    return cellFocusDetails(
+      cellId,
+      get(notebookAtom),
+      get(cellExecutionInputsAtom),
+    );
   });
 }
 
-function cellFocusDetails(cellId: CellId, notebookState: NotebookState) {
+function cellFocusDetails(
+  cellId: CellId,
+  notebookState: NotebookState,
+  executions: CellExecutionInputs,
+) {
   const { cellData, cellHandles, cellRuntime } = notebookState;
   const data = cellData[cellId];
   const runtime = cellRuntime[cellId];
@@ -107,12 +126,14 @@ function cellFocusDetails(cellId: CellId, notebookState: NotebookState) {
     return null;
   }
   const getEditorView = () => handle.editorViewOrNull;
+  const execution = cellExecutionView(executions, cellId);
 
   return {
     cellId,
     name: data.name,
     config: data.config,
-    status: runtime ? runtime.status : "idle",
+    status: execution.status,
+    blocked: execution.blocked,
     getEditorView: getEditorView,
     hasOutput: runtime?.output != null,
     hasConsoleOutput: runtime?.consoleOutputs != null,

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-8404): Outlines come from the host.
+ */
+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];
+}

package/src/core/cells/outputs.ts

@@ -1,15 +1,15 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useAtomValue } from "jotai";
-import { selectAtom } from "jotai/utils";
-import { useEffect, useMemo, useState } from "react";
+import { useEffect, useState } from "react";
 import { useEvent } from "@/hooks/useEvent";
-import { Logger } from "@/utils/Logger";
+import {
+  useNotebookCellConfig,
+  useNotebookDocumentEdit,
+} from "../notebook/document";
 import type { OutputMessage } from "../kernel/messages";
 import { viewStateAtom } from "../mode";
-import { getRequestClient } from "../network/requests";
-import { cellDataAtom, useCellActions } from "./cells";
-import type { CellId } from "./ids";
+import { type CellId, SCRATCH_CELL_ID } from "./ids";
 
 // This does not need to be overcomplicated. We can just store the expanded
 // state in a global map instead of Jotai since state is not shared between cells.
@@ -27,20 +27,9 @@ const expandedConsoleOutputs: Record<CellId, boolean> = {};
  */
 export function useExpandedOutput(cellId: CellId) {
   const { mode } = useAtomValue(viewStateAtom);
-  const { updateCellConfig } = useCellActions();
-  // Select just the flag: every output would otherwise re-render on each
-  // keystroke in its cell.
-  const isConfigExpanded = useAtomValue(
-    useMemo(
-      () =>
-        selectAtom(
-          cellDataAtom(cellId),
-          (data) => data?.config.expand_output ?? false,
-          Object.is,
-        ),
-      [cellId],
-    ),
-  );
+  const edit = useNotebookDocumentEdit();
+  const isConfigExpanded =
+    useNotebookCellConfig(cellId)?.expand_output ?? false;
   const [isSessionExpanded, setIsSessionExpanded] = useState(
     () => expandedOutputs[cellId] ?? isConfigExpanded,
   );
@@ -55,16 +44,11 @@ export function useExpandedOutput(cellId: CellId) {
 
   const setIsExpanded = useEvent((expanded: boolean) => {
     setIsSessionExpanded(expanded);
-    if (!isEditable) {
+    // The scratchpad is not in the document.
+    if (!isEditable || cellId === SCRATCH_CELL_ID) {
       return;
     }
-    const config = { expand_output: expanded };
-    updateCellConfig({ cellId, config });
-    getRequestClient()
-      .saveCellConfig({ configs: { [cellId]: config } })
-      .catch((error) => {
-        Logger.error("Failed to save expand_output config", error);
-      });
+    edit((tx) => tx.updateConfig(cellId, { expand_output: expanded }));
   });
 
   return [isExpanded, setIsExpanded] as const;