@marimo-team/islands 0.25.3-dev17 → 0.25.3-dev173

package/src/core/cells/__tests__/session.test.ts

@@ -46,10 +46,17 @@ describe("notebookStateFromSession", () => {
 
   const createNotebookCell = (
     id: string,
-    code: string | null = null,
-    name: string | null = null,
-    config: NotebookCell["config"] | null = null,
-    code_hash: string | null = null,
+    {
+      code = null,
+      name = null,
+      config = null,
+      code_hash = null,
+    }: {
+      code?: string | null;
+      name?: string | null;
+      config?: NotebookCell["config"] | null;
+      code_hash?: string | null;
+    } = {},
   ): NotebookCell => ({
     id,
     code,
@@ -87,7 +94,7 @@ describe("notebookStateFromSession", () => {
 
     it("logs error when session and notebook have different cells", () => {
       const session = createSession([createSessionCell("cell-1")]);
-      const notebook = createNotebook([createNotebookCell("cell-2")]);
+      const notebook = createNotebook([createNotebookCell("cell-2", {})]);
 
       const result = notebookStateFromSession(session, notebook);
 
@@ -328,7 +335,10 @@ describe("notebookStateFromSession", () => {
   describe("notebook only scenarios", () => {
     it("creates state from notebook with single cell", () => {
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "print('hello')", "my_cell"),
+        createNotebookCell("cell-1", {
+          code: "print('hello')",
+          name: "my_cell",
+        }),
       ]);
       const result = notebookStateFromSession(null, notebook);
 
@@ -356,11 +366,15 @@ describe("notebookStateFromSession", () => {
 
     it("creates state from notebook with custom config", () => {
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "code", "name", {
-          hide_code: true,
-          disabled: true,
-          column: 1,
-          expand_output: true,
+        createNotebookCell("cell-1", {
+          code: "code",
+          name: "name",
+          config: {
+            hide_code: true,
+            disabled: true,
+            column: 1,
+            expand_output: true,
+          },
         }),
       ]);
       const result = notebookStateFromSession(null, notebook);
@@ -464,7 +478,10 @@ describe("notebookStateFromSession", () => {
         ),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "print('hello')", "my_cell"),
+        createNotebookCell("cell-1", {
+          code: "print('hello')",
+          name: "my_cell",
+        }),
       ]);
       const result = notebookStateFromSession(session, notebook);
 
@@ -495,8 +512,8 @@ describe("notebookStateFromSession", () => {
         createSessionCell("cell-2"),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-2"),
-        createNotebookCell("cell-1"),
+        createNotebookCell("cell-2", {}),
+        createNotebookCell("cell-1", {}),
       ]);
 
       const result = notebookStateFromSession(session, notebook);
@@ -514,8 +531,8 @@ describe("notebookStateFromSession", () => {
         createSessionCell("cell-2"),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-1"),
-        createNotebookCell("cell-2"),
+        createNotebookCell("cell-1", {}),
+        createNotebookCell("cell-2", {}),
       ]);
 
       const result = notebookStateFromSession(session, notebook);
@@ -669,12 +686,42 @@ describe("notebookStateFromSession", () => {
       };
       // Create notebook cells with code values 'a' through 'f'
       const notebookCells = [
-        createNotebookCell("cell-a", "a", null, null, hashes.a),
-        createNotebookCell("cell-b", "b", null, null, hashes.b),
-        createNotebookCell("cell-c", "c", null, null, hashes.c),
-        createNotebookCell("cell-d", "d", null, null, hashes.d),
-        createNotebookCell("cell-e", "e", null, null, hashes.e),
-        createNotebookCell("cell-f", "f", null, null, hashes.f),
+        createNotebookCell("cell-a", {
+          code: "a",
+          name: null,
+          config: null,
+          code_hash: hashes.a,
+        }),
+        createNotebookCell("cell-b", {
+          code: "b",
+          name: null,
+          config: null,
+          code_hash: hashes.b,
+        }),
+        createNotebookCell("cell-c", {
+          code: "c",
+          name: null,
+          config: null,
+          code_hash: hashes.c,
+        }),
+        createNotebookCell("cell-d", {
+          code: "d",
+          name: null,
+          config: null,
+          code_hash: hashes.d,
+        }),
+        createNotebookCell("cell-e", {
+          code: "e",
+          name: null,
+          config: null,
+          code_hash: hashes.e,
+        }),
+        createNotebookCell("cell-f", {
+          code: "f",
+          name: null,
+          config: null,
+          code_hash: hashes.f,
+        }),
       ];
       // Create session cells with matching code_hashes
       const sessionCells = [
@@ -787,11 +834,36 @@ describe("notebookStateFromSession", () => {
 
       // Notebook has cells with code 'a', 'c', 'z', 'e', 'g'
       const notebookCells = [
-        createNotebookCell("cell-a", "a", null, null, hashes.a),
-        createNotebookCell("cell-c", "c", null, null, hashes.c),
-        createNotebookCell("cell-z", "z", null, null, hashes.z),
-        createNotebookCell("cell-e", "e", null, null, hashes.e),
-        createNotebookCell("cell-g", "g", null, null, hashes.g),
+        createNotebookCell("cell-a", {
+          code: "a",
+          name: null,
+          config: null,
+          code_hash: hashes.a,
+        }),
+        createNotebookCell("cell-c", {
+          code: "c",
+          name: null,
+          config: null,
+          code_hash: hashes.c,
+        }),
+        createNotebookCell("cell-z", {
+          code: "z",
+          name: null,
+          config: null,
+          code_hash: hashes.z,
+        }),
+        createNotebookCell("cell-e", {
+          code: "e",
+          name: null,
+          config: null,
+          code_hash: hashes.e,
+        }),
+        createNotebookCell("cell-g", {
+          code: "g",
+          name: null,
+          config: null,
+          code_hash: hashes.g,
+        }),
       ];
 
       const session = createSession(sessionCells);
@@ -908,15 +980,24 @@ describe("notebookStateFromSession", () => {
 
       // Notebook has cells: cell-1 (1 / 0, no hash), cell-2 (mo.md("Hello"), hash=moMd), cell-3 (slider, hash=slider)
       const notebookCells = [
-        createNotebookCell("cell-1", "1 / 0", null, null, null), // no hash
-        createNotebookCell("cell-2", 'mo.md("Hello")', null, null, hashes.moMd),
-        createNotebookCell(
-          "cell-3",
-          "x = mo.ui.slider(0, 10)",
-          null,
-          null,
-          hashes.slider,
-        ),
+        createNotebookCell("cell-1", {
+          code: "1 / 0",
+          name: null,
+          config: null,
+          code_hash: null,
+        }), // no hash
+        createNotebookCell("cell-2", {
+          code: 'mo.md("Hello")',
+          name: null,
+          config: null,
+          code_hash: hashes.moMd,
+        }),
+        createNotebookCell("cell-3", {
+          code: "x = mo.ui.slider(0, 10)",
+          name: null,
+          config: null,
+          code_hash: hashes.slider,
+        }),
       ];
 
       const session = createSession(sessionCells);

package/src/core/cells/__tests__/with-document-save.test.ts

@@ -0,0 +1,67 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { createDocumentModel } from "@/core/notebook/document";
+import { document } from "@/core/notebook/document/__tests__/fixtures";
+import type { Document } from "@/core/notebook/replica/resources/document";
+import {
+  type FakeEdit,
+  fakeBackend,
+} from "@/core/notebook/replica/__tests__/fixtures";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { store } from "@/core/state/jotai";
+import { Deferred } from "@/utils/Deferred";
+import { withDocumentSave } from "../document";
+
+const lifetimes: AbortController[] = [];
+
+afterEach(() => {
+  lifetimes.splice(0).forEach((controller) => controller.abort());
+  store.set(activeDocumentModelAtom, null);
+});
+
+describe("withDocumentSave", () => {
+  it("sends an edit made during the save only after the save finishes", async () => {
+    const source = atom<Document | null>(document());
+    const base = document();
+    const accepted: Document = {
+      ...base,
+      revision: { ...base.revision, version: 2, etag: "document-2" },
+    };
+    // Answer the edit, then publish the accepted document as the host would.
+    const edit = vi.fn<FakeEdit>().mockImplementation(async () => {
+      queueMicrotask(() => store.set(source, accepted));
+      return {
+        ok: true,
+        status: 200,
+        data: { document: accepted.revision, cells: {} },
+      };
+    });
+    const model = createDocumentModel({
+      store,
+      documentAtom: source,
+      backendAtom: atom(fakeBackend({ edit })),
+    });
+    const controller = new AbortController();
+    lifetimes.push(controller);
+    model.start({ signal: controller.signal });
+    store.set(activeDocumentModelAtom, model);
+
+    const saving = new Deferred<void>();
+    let flushing: Promise<unknown> | undefined;
+    const saved = withDocumentSave(async () => {
+      store.set(model.editAtom, (tx) => tx.setCode("c1", "x = 2"));
+      flushing = model.flush();
+      await Promise.resolve();
+      expect(edit).not.toHaveBeenCalled();
+      await saving.promise;
+      return "saved";
+    });
+
+    saving.resolve();
+    await expect(saved).resolves.toBe("saved");
+    await flushing;
+    expect(edit).toHaveBeenCalledTimes(1);
+  });
+});

package/src/core/cells/cell.ts

@@ -7,7 +7,12 @@ import { parseOutline } from "../dom/outline";
 import type { CellMessage, OutputMessage } from "../kernel/messages";
 import type { RuntimeState } from "../network/types";
 import { collapseConsoleOutputs } from "./collapseConsoleOutputs";
-import type { CellRuntimeState } from "./types";
+import type { CellRuntimeState, ConsoleOutputMessage } from "./types";
+
+// Ids increase with position, so the last output holds the maximum.
+function nextConsoleOutputId(outputs: ConsoleOutputMessage[]): number {
+  return (outputs.at(-1)?.id ?? 0) + 1;
+}
 
 export function transitionCell(
   cell: CellRuntimeState,
@@ -133,6 +138,7 @@ export function transitionCell(
           mimetype: output.mimetype,
           data: output.data,
           timestamp: output.timestamp,
+          id: output.id,
           response: output.response ?? "",
         };
       }
@@ -146,8 +152,13 @@ export function transitionCell(
     // messageConsole is an output that needs to be appended to the
     // existing console outputs.
     consoleOutputs = Array.isArray(messageConsole)
-      ? collapseConsoleOutputs(messageConsole)
-      : collapseConsoleOutputs([...consoleOutputs, messageConsole]);
+      ? collapseConsoleOutputs(
+          messageConsole.map((output, i) => ({ ...output, id: i + 1 })),
+        )
+      : collapseConsoleOutputs([
+          ...consoleOutputs,
+          { ...messageConsole, id: nextConsoleOutputId(consoleOutputs) },
+        ]);
   }
   nextCell.consoleOutputs = consoleOutputs;
   // Derive outline from output

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 {
@@ -47,6 +44,7 @@ import {
   createCell,
   createCellConfig,
   createCellRuntimeState,
+  isEdited,
 } from "./types";
 import {
   canUndoDeletes,
@@ -939,158 +937,14 @@ 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.
-   *
-   * 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.
-   */
-  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;
+  /** Follow the host document's cells and order. */
+  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: (
@@ -1114,6 +968,7 @@ const {
           mimetype: stdinOutput.mimetype,
           data: stdinOutput.data,
           timestamp: stdinOutput.timestamp,
+          id: stdinOutput.id,
           // An empty response resolves the prompt without retaining the secret.
           response: stdinOutput.mimetype === "text/password" ? "" : response,
         };
@@ -1126,36 +981,44 @@ const {
       },
     });
   },
-  setCells: (state, cells: CellData[]) => {
-    const cellData = Object.fromEntries(cells.map((cell) => [cell.id, cell]));
-
-    // WASM has no server-side SessionView to replay outputs, so the
-    // snapshot hydrated by notebookStateFromSession is the only source.
+  /** Record what the kernel last ran for each cell. */
+  setLastRun: (
+    state,
+    action: {
+      lastCodeRun: Record<CellId, string>;
+      lastExecutionTime: Record<CellId, number>;
+    },
+  ) => {
+    const cellData = { ...state.cellData };
+    for (const cellId of state.cellIds.inOrderIds) {
+      const cell = state.cellData[cellId];
+      const lastCodeRun = action.lastCodeRun[cellId] ?? null;
+      cellData[cellId] = {
+        ...cell,
+        lastCodeRun,
+        lastExecutionTime: action.lastExecutionTime[cellId] ?? null,
+        // A cell the kernel never ran is not stale on load.
+        edited: lastCodeRun !== null && isEdited(lastCodeRun, cell.code),
+      };
+    }
+    return { ...state, cellData };
+  },
+  /**
+   * Start each cell's runtime over for a new kernel. WASM keeps the outputs
+   * its mount-time snapshot gave it.
+   */
+  resetCellRuntimes: (state) => {
     const preserveSnapshot = isWasm();
-    const runtimeFor = (cellId: CellId): CellRuntimeState => {
-      if (!preserveSnapshot) {
-        return createCellRuntimeState();
-      }
-      const prev = state.cellRuntime[cellId];
+    const cellRuntime = { ...state.cellRuntime };
+    for (const cellId of [...state.cellIds.inOrderIds, SCRATCH_CELL_ID]) {
+      const previous = state.cellRuntime[cellId];
       const hasSnapshot =
-        prev && (prev.output != null || prev.consoleOutputs.length > 0);
-      return hasSnapshot ? prev : createCellRuntimeState();
-    };
-    const cellRuntime = Object.fromEntries(
-      cells.map((cell) => [cell.id, runtimeFor(cell.id)]),
-    );
-
-    return withScratchCell({
-      ...state,
-      cellIds: MultiColumn.fromIdsAndColumns(
-        cells.map((cell) => [cell.id, cell.config.column]),
-      ),
-      cellData: cellData,
-      cellRuntime: cellRuntime,
-      cellHandles: Object.fromEntries(
-        cells.map((cell) => [cell.id, createRef()]),
-      ),
-    });
+        preserveSnapshot &&
+        previous !== undefined &&
+        (previous.output != null || previous.consoleOutputs.length > 0);
+      cellRuntime[cellId] = hasSnapshot ? previous : createCellRuntimeState();
+    }
+    return { ...state, cellRuntime };
   },
   /**
    * Move focus to next cell
@@ -1439,7 +1302,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();
 
@@ -1487,7 +1357,11 @@ const {
       return state;
     }
 
-    updateEditorCodeFromPython(cellHandle.editorView, snapshot);
+    const update = prepareCellCodeUpdate(cellHandle.editorView.state, snapshot);
+
+    if (update) {
+      cellHandle.editorView.dispatch(update);
+    }
 
     return {
       ...state,
@@ -1598,9 +1472,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 (
@@ -1761,7 +1635,8 @@ export const cellErrorsAtom = atom((get) => {
 });
 
 export const notebookOutline = atom((get) => {
-  const { cellIds, cellRuntime } = get(notebookAtom);
+  const cellIds = get(cellIdsAtom);
+  const cellRuntime = get(cellsRuntimeAtom);
   const outlines = cellIds.inOrderIds.map(
     (cellId) => cellRuntime[cellId].outline,
   );