@marimo-team/islands 0.25.3-dev169 → 0.25.3-dev170

package/src/core/cells/cells.ts

@@ -1,11 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { historyField } from "@codemirror/commands";
-import { dequal as isEqual } from "dequal";
 import { type Atom, atom, useAtom, useAtomValue } from "jotai";
 import { atomFamily, selectAtom, splitAtom } from "jotai/utils";
 import { createRef, type ReducerWithoutAction } from "react";
 import type { CellHandle } from "@/components/editor/notebook-cell";
+import type { EditableDocument } from "@/core/notebook/document";
 import type { CollapsibleTree } from "@/utils/id-tree";
 import {
   type CellColumnId,
@@ -20,20 +20,17 @@ import { extractAllTracebackInfo, type TracebackInfo } from "@/utils/traceback";
 import { createReducerAndAtoms } from "../../utils/createReducer";
 import { foldAllBulk, unfoldAllBulk } from "../codemirror/editing/commands";
 import { editorMountScheduler } from "../codemirror/editor-mount-scheduler";
-import {
-  splitEditor,
-  updateEditorCodeFromPython,
-} from "../codemirror/language/utils";
+import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
+import { splitEditor } from "../codemirror/language/utils";
 import type { SerializedEditorState } from "../codemirror/types";
 import { findCollapseRange, mergeOutlines } from "../dom/outline";
 import type { CellMessage } from "../kernel/messages";
 import { isErrorMime } from "../mime";
 import type { CellConfig } from "../network/types";
-import { canUseRtc, isRtcEnabled } from "../rtc/state";
 import { createDeepEqualAtom, store } from "../state/jotai";
 import { isWasm } from "../wasm/utils";
 import { prepareCellForExecution, transitionCell } from "./cell";
-import { documentTransactionMiddleware } from "./document-changes";
+import { applyDocument, documentEditMiddleware } from "./document";
 import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
 import { type CellLog, getCellLogsForMessage } from "./logs";
 import {
@@ -939,158 +936,19 @@ const {
       cellLogs: [...state.cellLogs, ...action.logs],
     };
   },
-  setCellIds: (state, action: { cellIds: CellId[] }) => {
-    const isTheSame = isEqual(state.cellIds.inOrderIds, action.cellIds);
-    if (isTheSame) {
-      return state;
-    }
-
-    // Create new cell data and runtime states for the new cell IDs
-    const nextCellData = { ...state.cellData };
-    const nextCellRuntime = { ...state.cellRuntime };
-    const nextCellHandles = { ...state.cellHandles };
-
-    for (const cellId of action.cellIds) {
-      if (!(cellId in state.cellData)) {
-        nextCellData[cellId] = createCell({ id: cellId });
-      }
-      if (!(cellId in state.cellRuntime)) {
-        nextCellRuntime[cellId] = createCellRuntimeState();
-      }
-      if (!(cellId in state.cellHandles)) {
-        nextCellHandles[cellId] = createRef();
-      }
-    }
-
-    return {
-      ...state,
-      cellIds: MultiColumn.fromWithPreviousShape(action.cellIds, state.cellIds),
-      cellData: nextCellData,
-      cellRuntime: nextCellRuntime,
-      cellHandles: nextCellHandles,
-    };
-  },
   /**
-   * Rebuild the MultiColumn tree using each cell's `config.column` value.
+   * Follow the host document's cells and order.
    *
-   * Used after a transaction whose `set-config` changes updated cells'
-   * column metadata without physically moving them in the tree. Cells with
-   * `config.column == null` inherit the column of the previous cell in the
-   * given order (see `MultiColumn.fromIdsAndColumns`), which lets the server
-   * send column changes only at column boundaries.
+   * Dispatched without middleware, since the document already has these
+   * changes.
    */
-  rebuildCellColumns: (state, action: { cellIds: CellId[] }) => {
-    const newCellIds = MultiColumn.fromIdsAndColumns(
-      action.cellIds.map((id) => [
-        id,
-        state.cellData[id]?.config.column ?? null,
-      ]),
-    );
-    return { ...state, cellIds: newCellIds };
-  },
-  setCellCodes: (
-    state,
-    action: {
-      codes: string[];
-      ids: CellId[];
-      codeIsStale: boolean;
-      names?: string[];
-      configs?: CellConfig[];
-    },
-  ) => {
-    invariant(
-      action.codes.length === action.ids.length,
-      "Expected codes and ids to have the same length",
-    );
-
-    let nextState = { ...state };
-
-    const cellReducer = ({
-      cell,
-      code,
-      cellId,
-      name,
-      config,
-    }: {
-      cell: CellData | undefined;
-      code: string;
-      cellId: CellId;
-      name?: string;
-      config?: CellConfig;
-    }) => {
-      if (!cell) {
-        return createCell({
-          id: cellId,
-          code,
-          name: name,
-          config: config,
-          lastCodeRun: action.codeIsStale ? null : code,
-          edited: action.codeIsStale && code.trim().length > 0,
-        });
-      }
-
-      // If code is stale, we don't promote it to lastCodeRun
-      const lastCodeRun = action.codeIsStale ? cell.lastCodeRun : code;
-
-      // Mark as edited if the code has changed
-      const edited = lastCodeRun
-        ? lastCodeRun.trim() !== code.trim()
-        : Boolean(code);
-
-      // No change
-      if (cell.code.trim() === code.trim()) {
-        return {
-          ...cell,
-          code: code,
-          edited,
-          lastCodeRun,
-          ...(name !== undefined && { name }),
-          ...(config !== undefined && { config }),
-        };
-      }
-
-      // Update codemirror if mounted. RTC-backed cells sync via Loro; others
-      // (e.g. scratchpad) still need an explicit editor update.
-      if (!isRtcEnabled() || !canUseRtc(cellId)) {
-        const cellHandle = nextState.cellHandles[cellId].current;
-        if (cellHandle?.editorViewOrNull) {
-          updateEditorCodeFromPython(cellHandle.editorViewOrNull, code);
-        }
-      }
-
-      return {
-        ...cell,
-        code: code,
-        edited,
-        lastCodeRun,
-        ...(name !== undefined && { name }),
-        ...(config !== undefined && { config }),
-      };
-    };
-
-    for (let i = 0; i < action.ids.length; i++) {
-      const cellId = action.ids[i];
-      const code = action.codes[i];
-      if (cellId === undefined || code === undefined) {
-        continue;
+  reconcileDocument: (state, document: EditableDocument) => {
+    const nextState = applyDocument(state, document);
+    for (const cellId of state.cellIds.inOrderIds) {
+      if (!(cellId in nextState.cellData)) {
+        releaseCellAtoms(cellId);
       }
-      const name = action.names?.[i];
-      const config = action.configs?.[i];
-      nextState = {
-        ...nextState,
-        cellData: {
-          ...nextState.cellData,
-          [cellId]: cellReducer({
-            cell: nextState.cellData[cellId],
-            code,
-            cellId,
-            name,
-            config,
-          }),
-        },
-      };
     }
-
     return nextState;
   },
   setStdinResponse: (
@@ -1440,7 +1298,14 @@ const {
       cellHandle.editorView,
     );
 
-    updateEditorCodeFromPython(cellHandle.editorView, beforeCursorCode);
+    const update = prepareCellCodeUpdate(
+      cellHandle.editorView.state,
+      beforeCursorCode,
+    );
+
+    if (update) {
+      cellHandle.editorView.dispatch(update);
+    }
 
     const newCellId = CellId.create();
 
@@ -1488,7 +1353,11 @@ const {
       return state;
     }
 
-    updateEditorCodeFromPython(cellHandle.editorView, snapshot);
+    const update = prepareCellCodeUpdate(cellHandle.editorView.state, snapshot);
+
+    if (update) {
+      cellHandle.editorView.dispatch(update);
+    }
 
     return {
       ...state,
@@ -1599,9 +1468,9 @@ const {
 });
 
 // We apply the middleware here (rather than inline in createReducerAndAtoms)
-// so that the document transaction middleware can import CellActions and
+// so that the document edit middleware can import CellActions and
 // strictly type the dispatched actions without creating a circular dependency.
-addMiddleware(documentTransactionMiddleware);
+addMiddleware(documentEditMiddleware);
 
 function isCellCodeHidden(state: NotebookState, cellId: CellId): boolean {
   return (

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,360 @@
+/* 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 type * as Protocol from "@marimo-team/client/protocol";
+import { dequal as isEqual } from "dequal";
+import { createRef, type RefObject } from "react";
+import type { CellHandle } from "@/components/editor/notebook-cell";
+import type {
+  DocumentEdit,
+  DocumentEditor,
+  EditableCell,
+  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 { CellActions, NotebookState } from "./cells";
+import { groupByColumn } from "./columns";
+import type { CellId } from "./ids";
+import {
+  type CellData,
+  createCell,
+  createCellConfig,
+  createCellRuntimeState,
+  isEdited,
+} from "./types";
+
+type CellAction = DispatchedActionOf<CellActions>;
+
+// Following the document ----------------------------------------------------
+
+/** Rebuild the tree for `columns`, keeping column ids and collapsed sections. */
+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.fromWithPreviousShape(ids, previous[index]),
+    ),
+  );
+}
+
+/** 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 {
+  const cellIds = reconcileColumns(state.cellIds, groupByColumn(document));
+  let { cellData, cellRuntime, cellHandles } = state;
+  for (const id of cellIds.inOrderIds) {
+    const cell = document.cells.get(id);
+    const current = state.cellData[id];
+    if (cell === undefined) {
+      continue;
+    }
+    if (current === undefined) {
+      // A new cell has no run to be edited from.
+      cellData = {
+        ...cellData,
+        [id]: createCell({
+          id,
+          code: cell.code ?? "",
+          name: cell.name,
+          config: createCellConfig(cell.config),
+        }),
+      };
+      cellRuntime = { ...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.inOrderIds);
+  const removed = state.cellIds.inOrderIds.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 };
+}
+
+// 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.
+ */
+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) {
+    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) {
+      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;
+  }
+  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): Each action family that moves onto the model removes its
+ * part of `diffNotebook`. The middleware goes with the last one, code.
+ */
+export function documentEditMiddleware(
+  prev: NotebookState,
+  next: NotebookState,
+  action: CellAction,
+): void {
+  // The kernel's cells replace the notebook rather than edit it.
+  if (action.type === "setCells") {
+    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/types.ts

@@ -8,6 +8,18 @@ import type { CellId } from "./ids";
 import { DEFAULT_CELL_NAME } from "./names";
 import type { Outline } from "./outline";
 
+/**
+ * Whether a cell's code has changed since it last ran. A cell that never ran
+ * is edited once it has code. Whitespace at the ends does not count.
+ *
+ * TODO(MO-8028): Temporary until staleness comes from the host's executions.
+ */
+export function isEdited(lastCodeRun: string | null, code: string): boolean {
+  return lastCodeRun === null
+    ? code.trim().length > 0
+    : lastCodeRun.trim() !== code.trim();
+}
+
 /**
  * Create a new cell with default state.
  */

package/src/core/codemirror/cells/__tests__/extensions.test.ts

@@ -11,7 +11,8 @@ import { formattingChangeEffect } from "../../format";
 import { loroSyncAnnotation } from "../../rtc/loro/sync";
 import { exportedForTesting } from "../extensions";
 
-const { shouldAutorunMarkdownUpdate } = exportedForTesting;
+const { isCollaboratorUpdate, shouldAutorunMarkdownUpdate } =
+  exportedForTesting;
 
 function createTransaction(spec: TransactionSpec) {
   return EditorState.create({ doc: "" }).update(spec);
@@ -125,3 +126,28 @@ describe("shouldAutorunMarkdownUpdate", () => {
     ).toBe(false);
   });
 });
+
+describe("isCollaboratorUpdate", () => {
+  it("recognizes an update made only of Loro sync transactions", () => {
+    const transaction = createTransaction({
+      changes: { from: 0, insert: "#" },
+      annotations: [loroSyncAnnotation.of(true)],
+    });
+
+    expect(isCollaboratorUpdate([transaction])).toBe(true);
+  });
+
+  it("treats typing beside a Loro sync transaction as local", () => {
+    const remote = createTransaction({
+      changes: { from: 0, insert: "#" },
+      annotations: [loroSyncAnnotation.of(true)],
+    });
+    const local = createTransaction({
+      changes: { from: 0, insert: "x" },
+      annotations: [Transaction.userEvent.of("input.type")],
+    });
+
+    expect(isCollaboratorUpdate([remote, local])).toBe(false);
+    expect(isCollaboratorUpdate([])).toBe(false);
+  });
+});