@marimo-team/islands 0.25.3-dev18 → 0.25.3-dev181

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,392 @@
+/* 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 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 { 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 {
+  type CellData,
+  createCell,
+  createCellConfig,
+  createCellRuntimeState,
+  isEdited,
+} from "./types";
+
+type CellAction = DispatchedActionOf<ReducerCellActions>;
+
+// 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 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]: 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 };
+}
+
+/**
+ * Reconcile the model's document 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);
+    }
+  };
+  const unsubscribe = store.sub(model.documentAtom, follow);
+  signal.addEventListener("abort", unsubscribe, { once: true });
+  follow();
+}
+
+// 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 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/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;

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

@@ -134,12 +134,8 @@ export function usePendingDeleteService() {
       submit({ cellIds, deleteCell, deleteManyCells, store });
     },
     clear,
-    get idle() {
-      return entries.size === 0;
-    },
-    get shouldConfirmDelete() {
-      return entries.size > 1;
-    },
+    idle: entries.size === 0,
+    shouldConfirmDelete: entries.size > 1,
   };
 }
 

package/src/core/cells/session.ts

@@ -56,7 +56,7 @@ function mergeSessionAndNotebookCells(
     return {
       cellIds,
       sessionCellData: new Map(
-        cellIds.map((id, idx) => [id, null] as [CellId, SessionCell | null]),
+        cellIds.map((id) => [id, null] as [CellId, SessionCell | null]),
       ),
       notebookCellData: new Map(
         cellIds.map(

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.
  */
@@ -95,7 +107,7 @@ export interface CellRuntimeState {
   /** TOC outline */
   outline: Outline | null;
   /** messages encoding the cell's console outputs. */
-  consoleOutputs: WithResponse<OutputMessage>[];
+  consoleOutputs: ConsoleOutputMessage[];
   /** current status of the cell */
   status: RuntimeState;
   /** whether the cell has stale inputs*/
@@ -130,3 +142,12 @@ export type WithResponse<T> = T & {
    */
   response?: string;
 };
+
+/**
+ * A console output as held by the frontend. The id is assigned when the
+ * output arrives and never changes, so it can identify the output while
+ * neighboring outputs collapse, truncate, or receive a stdin response.
+ */
+export type ConsoleOutputMessage = WithResponse<OutputMessage> & {
+  id?: number;
+};

package/src/core/codemirror/__tests__/setup.test.ts

@@ -100,7 +100,7 @@ function getDuplicateKeymaps(state: EditorState) {
     (keymap) => keymap,
   );
   const duplicates = Objects.fromEntries(
-    Object.entries(groupBy).filter(([key, value]) => value.length > 1),
+    Object.entries(groupBy).filter(([_key, value]) => value.length > 1),
   );
   return duplicates;
 }
@@ -198,22 +198,22 @@ test("go to definition falls through to lower-priority LSP keymap", () => {
 test("placeholder adds another extension", () => {
   const opts = getOpts();
   const withAI = new PythonLanguageAdapter()
-    .getExtension(
-      opts.cellId,
-      opts.completionConfig,
-      opts.hotkeys,
-      "marimo-import",
-      opts.lspConfig,
-    )
+    .getExtension({
+      cellId: opts.cellId,
+      completionConfig: opts.completionConfig,
+      hotkeys: opts.hotkeys,
+      placeholderType: "marimo-import",
+      lspConfig: opts.lspConfig,
+    })
     .flat();
   const withoutAI = new PythonLanguageAdapter()
-    .getExtension(
-      opts.cellId,
-      opts.completionConfig,
-      opts.hotkeys,
-      "none",
-      opts.lspConfig,
-    )
+    .getExtension({
+      cellId: opts.cellId,
+      completionConfig: opts.completionConfig,
+      hotkeys: opts.hotkeys,
+      placeholderType: "none",
+      lspConfig: opts.lspConfig,
+    })
     .flat();
   expect(withAI.length - 1).toBe(withoutAI.length);
 });
@@ -221,22 +221,22 @@ test("placeholder adds another extension", () => {
 test("ai adds more extensions", () => {
   const opts = getOpts();
   const withAI = new PythonLanguageAdapter()
-    .getExtension(
-      opts.cellId,
-      opts.completionConfig,
-      opts.hotkeys,
-      "ai",
-      opts.lspConfig,
-    )
+    .getExtension({
+      cellId: opts.cellId,
+      completionConfig: opts.completionConfig,
+      hotkeys: opts.hotkeys,
+      placeholderType: "ai",
+      lspConfig: opts.lspConfig,
+    })
     .flat();
   const withoutAI = new PythonLanguageAdapter()
-    .getExtension(
-      opts.cellId,
-      opts.completionConfig,
-      opts.hotkeys,
-      "none",
-      opts.lspConfig,
-    )
+    .getExtension({
+      cellId: opts.cellId,
+      completionConfig: opts.completionConfig,
+      hotkeys: opts.hotkeys,
+      placeholderType: "none",
+      lspConfig: opts.lspConfig,
+    })
     .flat();
   expect(withAI.length - 2).toBe(withoutAI.length);
 });