@marimo-team/islands 0.25.3-dev209 → 0.25.3-dev210

package/src/core/cells/document.ts

@@ -1,8 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 /**
- * Keep the notebook state and the host document in step while cells live
- * in both.
+ * 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-8028): Delete this module once components read cells from the
  * document model and `notebookAtom` holds only runtime and editor state.
@@ -10,20 +11,15 @@
 
 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 { createRef } from "react";
 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 { store } from "../state/jotai";
 import { activeDocumentModelAtom } from "../session/active-notebook";
-import type { NotebookState, ReducerCellActions } from "./cells";
+import type { NotebookState } from "./cells";
 import type { CellId } from "./ids";
 import { followColumnsAtom } from "./layout";
 import {
@@ -34,16 +30,10 @@ import {
   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 {
+/** 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)
@@ -59,13 +49,6 @@ function followCell(
   if (code === current.code) {
     return { ...current, name: cell.name, config };
   }
-  // The editor takes the host's code; one that shows it already, as the
-  // editor this client types in does, is left alone.
-  const view = handle?.current?.editorViewOrNull;
-  const update = view && prepareCellCodeUpdate(view.state, code);
-  if (update) {
-    view.dispatch(update);
-  }
   return {
     ...current,
     code,
@@ -110,7 +93,7 @@ export function applyDocument(
           code,
           name: cell.name,
           config: createCellConfig(cell.config),
-          edited: code.trim().length > 0,
+          edited: isEdited(null, code),
         }),
       };
       cellRuntime = {
@@ -120,7 +103,7 @@ export function applyDocument(
       cellHandles = { ...cellHandles, [id]: createRef() };
       continue;
     }
-    const next = followCell(current, cell, state.cellHandles[id]);
+    const next = followCell(current, cell);
     if (next !== current) {
       cellData = { ...cellData, [id]: next };
     }
@@ -176,64 +159,6 @@ export function followDocument({
   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 -------------------------------
 
 /**

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

@@ -17,7 +17,6 @@ import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { requestClientAtom } from "@/core/network/requests";
 import type { MarimoConfig } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
-import { type CodemirrorCellActions, cellActionsState } from "../cells/state";
 import { cellIdState } from "../config/extension";
 import { formatAll, formatEditorViews, formatSQL } from "../format";
 import {
@@ -44,7 +43,6 @@ vi.mock("@/core/config/config", () => ({
   }),
 }));
 
-const updateCellCode = vi.fn();
 const createdViews: EditorView[] = [];
 
 function createEditor(content: string, cellId: CellId) {
@@ -65,9 +63,6 @@ function createEditor(content: string, cellId: CellId) {
         lspConfig: {},
       }),
       cellIdState.of(cellId),
-      cellActionsState.of({
-        updateCellCode,
-      } as unknown as CodemirrorCellActions),
     ],
   });
 
@@ -85,7 +80,6 @@ const mockConfig = {
 } as MarimoConfig;
 
 beforeEach(() => {
-  updateCellCode.mockClear();
   vi.clearAllMocks();
   // Set the mock request client in the atom
   store.set(requestClientAtom, mockRequestClient);
@@ -131,16 +125,6 @@ describe("format", () => {
 
       expect(views[cid1].state.doc.toString()).toBe(formattedCode1);
       expect(views[cid2].state.doc.toString()).toBe(formattedCode2);
-      expect(updateCellCode).toHaveBeenCalledWith({
-        cellId: cid1,
-        code: formattedCode1,
-        formattingChange: true,
-      });
-      expect(updateCellCode).toHaveBeenCalledWith({
-        cellId: cid2,
-        code: formattedCode2,
-        formattingChange: true,
-      });
     });
 
     it("should not update editor if formatted code is same as original", async () => {
@@ -161,7 +145,6 @@ describe("format", () => {
       await formatEditorViews(views);
 
       expect(views[cid].state.doc.toString()).toBe(originalCode);
-      expect(updateCellCode).not.toHaveBeenCalled();
     });
   });
 
@@ -194,12 +177,6 @@ describe("format", () => {
         },
         lineLength: 88,
       });
-
-      expect(updateCellCode).toHaveBeenCalledWith({
-        cellId: cid1,
-        code: "import numpy as np",
-        formattingChange: true,
-      });
     });
   });
 
@@ -220,11 +197,6 @@ describe("format", () => {
         WHERE
           id = 1"
       `);
-      expect(updateCellCode).toHaveBeenCalledWith({
-        cellId: cid,
-        code: editor.state.doc.toString(),
-        formattingChange: true,
-      });
     });
 
     it("should not format if language adapter is not SQL", async () => {
@@ -238,7 +210,6 @@ describe("format", () => {
       expect(editor.state.doc.toString()).toBe(
         "SELECT * FROM table WHERE id = 1",
       );
-      expect(updateCellCode).not.toHaveBeenCalled();
     });
   });
 
@@ -257,10 +228,5 @@ describe("format", () => {
       WHERE
         id = 1"
     `);
-    expect(updateCellCode).toHaveBeenCalledWith({
-      cellId: cid,
-      code: editor.state.doc.toString(),
-      formattingChange: true,
-    });
   });
 });

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

@@ -23,7 +23,7 @@ import { createObservable } from "@/core/state/observable";
 import { formatKeymapExtension } from "../extensions";
 import { formattingChangeEffect } from "../format";
 import { goToDefinitionAtCursorPosition } from "../go-to-definition/utils";
-import { getEditorCodeAsPython } from "../language/utils";
+import { getCellCode } from "../language/cell-code";
 import { isAtEndOfEditor, isAtStartOfEditor } from "../utils";
 import {
   breakpointGutter,
@@ -323,26 +323,28 @@ function cellKeymaps({
  */
 function cellCodeEditing(hotkeys: HotkeyProvider): Extension[] {
   const onChangePlugin = EditorView.updateListener.of((update) => {
-    if (update.docChanged) {
-      // Check if the doc update was a formatting change
-      // e.g. changing from python to markdown
-      const isFormattingChange = update.transactions.some((tr) =>
-        tr.effects.some((effect) => effect.is(formattingChangeEffect)),
-      );
-      const nextCode = getEditorCodeAsPython(update.view);
-      const cellActions = update.view.state.facet(cellActionsState);
-      const cellId = update.view.state.facet(cellIdState);
-      cellActions.updateCellCode({
+    if (!update.docChanged) {
+      return;
+    }
+    const cellId = update.view.state.facet(cellIdState);
+    // The code itself reaches the document through the editor's binding.
+    // Formatting, including a language switch, rewrites the code without
+    // editing it, so the cell stays as run or as edited as it was.
+    const formatted = update.transactions.some((tr) =>
+      tr.effects.some((effect) => effect.is(formattingChangeEffect)),
+    );
+    if (formatted) {
+      update.view.state.facet(cellActionsState).codeFormatted({
         cellId,
-        code: nextCode,
-        formattingChange: isFormattingChange,
+        before: getCellCode(update.startState),
+        after: getCellCode(update.state),
       });
+    }
 
-      // Clear pending cut state if this cell was marked for cut
-      const pendingCutCellIds = store.get(pendingCutCellIdsAtom);
-      if (pendingCutCellIds.has(cellId)) {
-        store.set(clearPendingCutAtom);
-      }
+    // Clear pending cut state if this cell was marked for cut
+    const pendingCutCellIds = store.get(pendingCutCellIdsAtom);
+    if (pendingCutCellIds.has(cellId)) {
+      store.set(clearPendingCutAtom);
     }
   });
 

package/src/core/codemirror/cells/state.ts

@@ -8,11 +8,6 @@ export interface CodemirrorCellActions extends CellActions {
   toggleHideCode: () => boolean;
   aiCellCompletion: () => boolean;
   createManyBelow: (content: string[]) => void;
-  /**
-   * Create the setup cell with the given code, or append the code to the
-   * existing setup cell if one already exists.
-   */
-  addOrAppendSetupCell: (code: string) => void;
   onRun: () => void;
   deleteCell: () => void;
   afterToggleMarkdown: () => void;

package/src/core/codemirror/format.ts

@@ -13,8 +13,6 @@ import { notebookCellEditorViews } from "../cells/utils";
 import { getResolvedMarimoConfig } from "../config/config";
 import type { ConnectionName } from "../datasets/engines";
 import { getRequestClient } from "../network/requests";
-import { cellActionsState } from "./cells/state";
-import { cellIdState } from "./config/extension";
 import { languageAdapterState } from "./language/extension";
 import { SCHEMA_CACHE } from "./language/languages/sql/completion-store";
 import { isKnownDialect } from "./language/languages/sql/utils";
@@ -54,15 +52,9 @@ export async function formatEditorViews(views: Record<CellId, EditorView>) {
       continue;
     }
 
-    const actions = view.state.facet(cellActionsState);
-    actions.updateCellCode({
-      cellId,
-      code: formattedCode,
-      formattingChange: true,
-    });
     const update = prepareCellCodeUpdate(view.state, formattedCode);
     if (update) {
-      view.dispatch(update);
+      view.dispatch(update, { effects: formattingChangeEffect.of(true) });
     }
   }
 }
@@ -111,16 +103,6 @@ export async function formatSQL(editor: EditorView, engine: ConnectionName) {
     return;
   }
 
-  // Update Python in the notebook state
-  const codeAsPython = languageAdapter.transformIn(formattedSQL)[0];
-  const actions = editor.state.facet(cellActionsState);
-  const cellId = editor.state.facet(cellIdState);
-  actions.updateCellCode({
-    cellId,
-    code: codeAsPython,
-    formattingChange: true,
-  });
-
   // Update editor with formatted SQL
   replaceEditorContent(editor, formattedSQL, {
     effects: [formattingChangeEffect.of(true)],

package/src/core/codemirror/misc/__tests__/paste.test.ts

@@ -443,7 +443,9 @@ def _():
     view.dom.onpaste = (evt) => extension[0].domEventHandlers.paste(evt, view);
     view.dom.dispatchEvent(event);
 
-    expect(addOrAppendSetupCell).toHaveBeenCalledWith("import numpy as np");
+    expect(addOrAppendSetupCell).toHaveBeenCalledWith({
+      code: "import numpy as np",
+    });
     expect(createManyBelow).toHaveBeenCalledWith(["x = 1"]);
   });
 });

package/src/core/codemirror/misc/paste.ts

@@ -23,7 +23,7 @@ export function pasteBundle(): Extension[] {
         // The `with app.setup` block belongs in the special setup cell, not in
         // an ordinary cell.
         if (setup !== null) {
-          actions.addOrAppendSetupCell(setup);
+          actions.addOrAppendSetupCell({ code: setup });
         }
         if (cells.length > 0) {
           actions.createManyBelow(cells);

package/src/core/notebook/document/__tests__/codemirror.test.ts

@@ -22,6 +22,10 @@ import {
   languageMetadataField,
   setLanguageMetadata,
 } from "../../../codemirror/language/metadata";
+import {
+  editorShowsCellCode,
+  getCellCode,
+} from "../../../codemirror/language/cell-code";
 import { SCHEMA_CACHE } from "../../../codemirror/language/languages/sql/completion-store";
 import type { Document } from "../../replica/resources/document";
 import { cellCodeBinding } from "../codemirror";
@@ -62,6 +66,7 @@ function languages(id: CellId) {
       aiCellCompletion: () => false,
       createManyBelow: vi.fn(),
       addOrAppendSetupCell: vi.fn(),
+      setCellCode: vi.fn(),
       createCell: vi.fn(),
       moveToNextCell: vi.fn(),
       addSetupCell: vi.fn(),
@@ -584,3 +589,170 @@ describe("cellCodeBinding before the cell has source", () => {
     expect(view.state.field(languageAdapterState).type).toBe("markdown");
   });
 });
+
+describe("cellCodeBinding under random interleavings", () => {
+  // A small deterministic generator, so a failure names the seed to replay.
+  function generator(seed: number) {
+    let state = seed >>> 0;
+    return () => {
+      state = (Math.imul(state, 1664525) + 1013904223) >>> 0;
+      return state / 2 ** 32;
+    };
+  }
+
+  const TOKENS = ["x", " ", "\n", "= 1", "mo.md(", ")", '"""', "#", "y + 2"];
+  const REMOTE = [
+    "x = 2",
+    'mo.md(r"""hi""")',
+    "y = x + 1\nz = y",
+    "",
+    "# other",
+  ];
+
+  async function stress(seed: number, adapters: boolean) {
+    const random = generator(seed);
+    const pick = <T>(items: readonly T[]) =>
+      items[Math.floor(random() * items.length)];
+    const { store, source, model, view, editor, code } = setup({
+      adapters,
+      initial: withCells(document(), { c1: cell("c1", "x = 1") }),
+    });
+    const other = editor();
+    await settle();
+    let composing = false;
+    Object.defineProperty(view, "composing", { get: () => composing });
+    let version = 1;
+    let hostCode = "x = 1";
+    const host = (next: string | null) => {
+      version += 1;
+      store.set(
+        source,
+        withCells(document(), {
+          c1:
+            next === null
+              ? null
+              : cell("c1", next, { version, etag: `c1-${version}` }),
+        }),
+      );
+    };
+    const agrees = (state = view.state) => {
+      const current = code();
+      return typeof current !== "string" || editorShowsCellCode(state, current);
+    };
+    // Every write to the model is one local or host step. A binding that
+    // echoed its own pulls back as pushes would write many more.
+    let writes = 0;
+    let steps = 0;
+    // Checks inside the loop collect here, so one assertion reports them all.
+    const failures: string[] = [];
+    const check = (label: string, ok: boolean) => {
+      if (!ok) {
+        failures.push(label);
+      }
+    };
+    const unsubscribe = store.sub(model.documentAtom, () => {
+      writes += 1;
+    });
+
+    for (let step = 0; step < 200; step++) {
+      const roll = random();
+      steps += 1;
+      if (roll < 0.35) {
+        const length = view.state.doc.length;
+        const from = Math.floor(random() * (length + 1));
+        view.dispatch({
+          changes: { from, insert: pick(TOKENS) },
+          userEvent: "input.type",
+        });
+        // A local change reaches the model before the update returns.
+        check(`step ${step} after typing`, agrees());
+      } else if (roll < 0.45) {
+        const length = view.state.doc.length;
+        if (length > 0) {
+          const from = Math.floor(random() * length);
+          const to = Math.min(length, from + 1 + Math.floor(random() * 5));
+          view.dispatch({
+            changes: { from, to },
+            userEvent: "delete.backward",
+          });
+          check(`step ${step} after deleting`, agrees());
+        }
+      } else if (roll < 0.6) {
+        // The host confirms what this client wrote.
+        const current = code();
+        if (typeof current === "string") {
+          hostCode = current;
+          host(hostCode);
+        }
+      } else if (roll < 0.75) {
+        // Another client changes the cell.
+        hostCode = pick(REMOTE);
+        host(hostCode);
+      } else if (roll < 0.8) {
+        // The cell leaves the document and returns.
+        host(null);
+        await settle();
+        host(hostCode);
+      } else if (roll < 0.88) {
+        composing = !composing;
+        if (!composing) {
+          view.contentDOM.dispatchEvent(new Event("compositionend"));
+        }
+      } else if (roll < 0.94) {
+        other.dispatch({
+          changes: { from: other.state.doc.length, insert: pick(TOKENS) },
+          userEvent: "input.type",
+        });
+      } else if (roll < 0.97) {
+        // The host withholds the cell's source for a while.
+        version += 1;
+        store.set(
+          source,
+          withCells(document(), {
+            c1: cell("c1", "", { code: null, version, etag: `c1-${version}` }),
+          }),
+        );
+      } else {
+        // The document goes away and comes back.
+        store.set(source, null);
+        await settle();
+        host(hostCode);
+      }
+      if (random() < 0.5) {
+        await settle();
+        check(`step ${step} after settling`, composing || agrees());
+        check(`step ${step} other editor`, agrees(other.state));
+      }
+    }
+    unsubscribe();
+    check(`${writes} writes for ${steps} steps`, writes <= 2 * steps);
+
+    composing = false;
+    view.contentDOM.dispatchEvent(new Event("compositionend"));
+    host(hostCode);
+    await settle();
+    await settle();
+    const current = code();
+    check("cell present at the end", typeof current === "string");
+    check("editor agrees at the end", agrees());
+    check("other editor agrees at the end", agrees(other.state));
+    // Without adapters the agreement is plain text equality.
+    check(
+      "text equal at the end",
+      adapters || getCellCode(view.state) === current,
+    );
+    return failures;
+  }
+
+  const seeds = Array.from({ length: 16 }, (_, i) => i + 1);
+  for (const seed of seeds) {
+    it(`keeps editor and model in agreement, seed ${seed}`, async () => {
+      expect(await stress(seed, false)).toEqual([]);
+    });
+  }
+  for (const seed of seeds) {
+    it(`keeps editor and model in agreement with language adapters, seed ${seed}`, async () => {
+      expect(await stress(seed, true)).toEqual([]);
+    });
+  }
+});

package/src/stories/layout/vertical/one-column.stories.tsx

@@ -3,11 +3,13 @@
 import type { Meta } from "@storybook/react-vite";
 import { createStore, Provider } from "jotai";
 import { createRef } from "react";
+import { mockDocumentModel } from "@/__mocks__/document";
 import { cellId } from "@/__tests__/branded";
 import { type NotebookState, notebookAtom } from "@/core/cells/cells";
 import { createCellRuntimeState } from "@/core/cells/types";
 import { defaultUserConfig, parseAppConfig } from "@/core/config/config-schema";
 import { showCodeInRunModeAtom } from "@/core/meta/state";
+import { DocumentProvider } from "@/core/notebook/document";
 import { connectionAtom } from "@/core/network/connection";
 import { requestClientAtom } from "@/core/network/requests";
 import { resolveRequestClient } from "@/core/network/resolve";
@@ -105,18 +107,21 @@ const EditModeCodeShown = () => {
   store.set(connectionAtom, { state: ConnectionState.OPEN });
   store.set(requestClientAtom, resolveRequestClient());
   store.set(showCodeInRunModeAtom, true);
+  const { model } = mockDocumentModel(store, null);
 
   return (
     <Provider store={store}>
-      <TooltipProvider>
-        <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
-          <CellArray
-            mode="edit"
-            userConfig={defaultUserConfig()}
-            appConfig={parseAppConfig({})}
-          />
-        </CellsRenderer>
-      </TooltipProvider>
+      <DocumentProvider model={model}>
+        <TooltipProvider>
+          <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
+            <CellArray
+              mode="edit"
+              userConfig={defaultUserConfig()}
+              appConfig={parseAppConfig({})}
+            />
+          </CellsRenderer>
+        </TooltipProvider>
+      </DocumentProvider>
     </Provider>
   );
 };
@@ -130,18 +135,21 @@ const EditModeCodeHidden = () => {
   store.set(connectionAtom, { state: ConnectionState.OPEN });
   store.set(requestClientAtom, resolveRequestClient());
   store.set(showCodeInRunModeAtom, true);
+  const { model } = mockDocumentModel(store, null);
 
   return (
     <Provider store={store}>
-      <TooltipProvider>
-        <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
-          <CellArray
-            mode="edit"
-            userConfig={defaultUserConfig()}
-            appConfig={parseAppConfig({})}
-          />
-        </CellsRenderer>
-      </TooltipProvider>
+      <DocumentProvider model={model}>
+        <TooltipProvider>
+          <CellsRenderer appConfig={parseAppConfig({})} mode="edit">
+            <CellArray
+              mode="edit"
+              userConfig={defaultUserConfig()}
+              appConfig={parseAppConfig({})}
+            />
+          </CellsRenderer>
+        </TooltipProvider>
+      </DocumentProvider>
     </Provider>
   );
 };