@marimo-team/frontend 0.25.3-dev102 → 0.25.3-dev104

package/src/core/document/__tests__/hooks.test.tsx

@@ -0,0 +1,386 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, render, renderHook } from "@testing-library/react";
+import { atom, createStore, Provider } from "jotai";
+import type { PropsWithChildren } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import type { Document } from "../../host/resources/document";
+import {
+  createDocumentModel,
+  DocumentProvider,
+  useCellCode,
+  useCellConfig,
+  useCellName,
+  useCellOrder,
+  useDocumentEdit,
+  useDocumentSync,
+  type DocumentHost,
+} from "..";
+import { cell, CONFIG, document, withCells } from "./fixtures";
+
+type HostEdit = DocumentHost["edit"];
+const lifetimes: AbortController[] = [];
+
+afterEach(() => {
+  act(() => lifetimes.splice(0).forEach((controller) => controller.abort()));
+});
+
+function setup(
+  initial: Document | null = document(),
+  edit: HostEdit = () => new Promise(() => {}),
+) {
+  const store = createStore();
+  const source = atom(initial);
+  const detached = vi.fn();
+  source.onMount = () => detached;
+  const controller = new AbortController();
+  lifetimes.push(controller);
+  const model = createDocumentModel({
+    store,
+    notebook: { documentAtom: source, edit },
+    signal: controller.signal,
+  });
+  const wrapper = ({ children }: PropsWithChildren) => (
+    <Provider store={store}>
+      <DocumentProvider model={model}>{children}</DocumentProvider>
+    </Provider>
+  );
+  return { store, source, model, controller, detached, wrapper };
+}
+
+describe("document hooks", () => {
+  it("flushes edits through the model and aborts work when the session ends", async () => {
+    const accepted = withCells(document(), { c1: cell("c1", "local") });
+    const edit = vi
+      .fn<HostEdit>()
+      .mockResolvedValueOnce({
+        ok: true,
+        status: 200,
+        data: { document: accepted.revision, cells: {} },
+      })
+      .mockImplementation(
+        () => new Deferred<Awaited<ReturnType<HostEdit>>>().promise,
+      );
+    const { store, source, wrapper, controller } = setup(document(), edit);
+    const { result, unmount } = renderHook(
+      () => ({
+        code: useCellCode("c1"),
+        edit: useDocumentEdit(),
+        sync: useDocumentSync(),
+      }),
+      { wrapper, reactStrictMode: true },
+    );
+    act(() => result.current.edit((tx) => tx.setCode("c1", "local")));
+    expect(result.current.sync.status).toEqual({ status: "pending" });
+    const { flush, retry } = result.current.sync;
+    await act(async () => {
+      const flushed = flush();
+      // Allow the response to arrive before publishing its confirmation.
+      await Promise.resolve();
+      store.set(source, accepted);
+      await expect(flushed).resolves.toEqual(accepted.revision);
+    });
+    expect(result.current.sync).toEqual({
+      status: { status: "idle" },
+      flush,
+      retry,
+    });
+    expect(result.current.code).toBe("local");
+    act(() => result.current.edit((tx) => tx.setCode("c1", "next")));
+    await act(async () => {
+      const rejected = (async () => {
+        await expect(flush()).rejects.toThrow("session ended");
+      })();
+      await Promise.resolve();
+      unmount();
+      controller.abort(new Error("session ended"));
+      await rejected;
+    });
+    expect(edit).toHaveBeenCalledTimes(2);
+    expect(edit.mock.calls[1][2].signal.aborted).toBe(true);
+  });
+
+  it("publishes compound edits atomically without changing host records", () => {
+    const initial = document();
+    const { store, source, wrapper } = setup(initial);
+    const seen = vi.fn();
+    const { result } = renderHook(
+      () => {
+        const value = {
+          code: useCellCode("c1"),
+          name: useCellName("c1"),
+          config: useCellConfig("c1"),
+          order: useCellOrder(),
+        };
+        seen(value);
+        return { ...value, edit: useDocumentEdit() };
+      },
+      { wrapper },
+    );
+    const edit = result.current.edit;
+    seen.mockClear();
+
+    act(() =>
+      edit((tx) => {
+        tx.setCode("c1", "edited");
+        tx.setName("c1", "renamed");
+        tx.updateConfig("c1", { hide_code: true });
+        tx.setOrder(["c2", "c1"]);
+      }),
+    );
+
+    const expected = {
+      code: "edited",
+      name: "renamed",
+      config: { ...CONFIG, hide_code: true },
+      order: ["c2", "c1"],
+    };
+    expect(result.current).toEqual({ ...expected, edit });
+    expect(seen.mock.calls.length).toBeGreaterThan(0);
+    for (const [value] of seen.mock.calls) {
+      expect(value).toEqual(expected);
+    }
+    expect(store.get(source)).toBe(initial);
+
+    expect(() =>
+      act(() =>
+        edit((tx) => {
+          tx.setCode("c1", "rolled back");
+          tx.deleteCell("missing");
+        }),
+      ),
+    ).toThrow("does not exist");
+    expect(result.current).toEqual({ ...expected, edit });
+  });
+
+  it("subscribes to individual fields and keeps the edit function stable", () => {
+    const { store, source, wrapper } = setup();
+    const onCode = vi.fn();
+    const onName = vi.fn();
+    const onConfig = vi.fn();
+    const onOrder = vi.fn();
+    const onEdit = vi.fn();
+    function Code() {
+      onCode(useCellCode("c1"));
+      return null;
+    }
+    function Name() {
+      onName(useCellName("c1"));
+      return null;
+    }
+    function Config() {
+      onConfig(useCellConfig("c1"));
+      return null;
+    }
+    function Order() {
+      onOrder(useCellOrder());
+      return null;
+    }
+    function Edit() {
+      onEdit(useDocumentEdit());
+      return null;
+    }
+    render(
+      <>
+        <Code />
+        <Name />
+        <Config />
+        <Order />
+        <Edit />
+      </>,
+      { wrapper },
+    );
+    const edit: ReturnType<typeof useDocumentEdit> = onEdit.mock.calls[0][0];
+    const clear = () =>
+      [onCode, onName, onConfig, onOrder, onEdit].forEach((fn) =>
+        fn.mockClear(),
+      );
+    clear();
+
+    act(() => edit((tx) => tx.setCode("c1", "local")));
+    expect(onCode).toHaveBeenCalledWith("local");
+    expect(onName).not.toHaveBeenCalled();
+    expect(onConfig).not.toHaveBeenCalled();
+    expect(onOrder).not.toHaveBeenCalled();
+    expect(onEdit).not.toHaveBeenCalled();
+    clear();
+
+    act(() => edit((tx) => tx.setCode("c2", "other")));
+    expect(
+      [onCode, onName, onConfig, onOrder, onEdit].map((fn) => fn.mock.calls),
+    ).toEqual([[], [], [], [], []]);
+
+    act(() =>
+      store.set(
+        source,
+        withCells(document(), {
+          c1: cell("c1", "remote", {
+            name: "host name",
+            config: { ...CONFIG },
+          }),
+        }),
+      ),
+    );
+    expect(onName).toHaveBeenCalledWith("host name");
+    expect(onCode).not.toHaveBeenCalled();
+    expect(onConfig).not.toHaveBeenCalled();
+    expect(onOrder).not.toHaveBeenCalled();
+    expect(onEdit).not.toHaveBeenCalled();
+    clear();
+
+    act(() =>
+      edit((tx) => {
+        tx.updateConfig("c1", { disabled: true });
+        tx.setOrder(["c2", "c1"]);
+      }),
+    );
+    expect(onConfig).toHaveBeenCalledWith({ ...CONFIG, disabled: true });
+    expect(onOrder).toHaveBeenCalledWith(["c2", "c1"]);
+    expect(onCode).not.toHaveBeenCalled();
+    expect(onName).not.toHaveBeenCalled();
+    expect(onEdit).not.toHaveBeenCalled();
+  });
+
+  it("switches field subscriptions when the cell id changes", () => {
+    const { wrapper } = setup();
+    const { result, rerender } = renderHook(
+      (id) => ({
+        code: useCellCode(id),
+        edit: useDocumentEdit(),
+      }),
+      { wrapper, initialProps: "c1" },
+    );
+    const edit = result.current.edit;
+    rerender("c2");
+    expect(result.current).toEqual({ code: "y = 2", edit });
+    act(() => edit((tx) => tx.setCode("c1", "old subscription")));
+    expect(result.current).toEqual({ code: "y = 2", edit });
+    act(() => edit((tx) => tx.setCode("c2", "new subscription")));
+    expect(result.current).toEqual({ code: "new subscription", edit });
+  });
+
+  it("distinguishes unavailable documents, hidden source, empty code, and deleted cells", () => {
+    const { store, source, wrapper } = setup(null);
+    const { result } = renderHook(
+      () => ({
+        order: useCellOrder(),
+        code: useCellCode("c1"),
+        name: useCellName("c1"),
+        config: useCellConfig("c1"),
+        edit: useDocumentEdit(),
+      }),
+      { wrapper },
+    );
+    const edit = result.current.edit;
+    const absent = {
+      code: undefined,
+      name: undefined,
+      config: undefined,
+      edit,
+    };
+    expect(result.current).toEqual({ ...absent, order: null });
+    const hidden = withCells(document(), {
+      c1: cell("c1", "", { code: null }),
+    });
+    act(() => store.set(source, hidden));
+    expect(result.current).toEqual({
+      code: null,
+      name: "_",
+      config: CONFIG,
+      order: ["c1", "c2"],
+      edit,
+    });
+    act(() => edit((tx) => tx.setCode("c1", "")));
+    expect(result.current.code).toBe("");
+    act(() => store.set(source, null));
+    expect(result.current).toEqual({ ...absent, order: null });
+    act(() => store.set(source, hidden));
+    expect(result.current.code).toBe("");
+    act(() =>
+      edit((tx) => {
+        tx.deleteCell("c1");
+        tx.deleteCell("c2");
+      }),
+    );
+    expect(result.current).toEqual({ ...absent, order: [] });
+  });
+
+  it("preserves drafts across Strict Mode remounts until the session ends", () => {
+    const { wrapper, controller, detached } = setup();
+    const useProbe = () => ({
+      code: useCellCode("c1"),
+      edit: useDocumentEdit(),
+    });
+    const first = renderHook(useProbe, { wrapper, reactStrictMode: true });
+    const edit = first.result.current.edit;
+    act(() => edit((tx) => tx.setCode("c1", "strict mode edit")));
+    expect(first.result.current.code).toBe("strict mode edit");
+    first.unmount();
+    expect(detached).not.toHaveBeenCalled();
+
+    const second = renderHook(useProbe, { wrapper, reactStrictMode: true });
+    expect(second.result.current.code).toBe("strict mode edit");
+    act(() => edit((tx) => tx.setCode("c1", "same session")));
+    expect(second.result.current.code).toBe("same session");
+    second.unmount();
+    controller.abort();
+    expect(detached).toHaveBeenCalledOnce();
+    const callback = vi.fn();
+    edit(callback);
+    expect(callback).not.toHaveBeenCalled();
+  });
+
+  it("switches reads and edits to the supplied model and store", () => {
+    const first = setup();
+    const second = setup(
+      withCells(document("n2"), {
+        c1: cell("c1", "second source", { notebook_id: "n2" }),
+      }),
+    );
+    const seen = vi.fn();
+    function Probe() {
+      seen({ code: useCellCode("c1"), edit: useDocumentEdit() });
+      return null;
+    }
+    const tree = ({ store, model }: typeof first) => (
+      <Provider store={store}>
+        <DocumentProvider model={model}>
+          <Probe />
+        </DocumentProvider>
+      </Provider>
+    );
+    const { rerender } = render(tree(first));
+    const oldEdit: ReturnType<typeof useDocumentEdit> =
+      seen.mock.calls.at(-1)![0].edit;
+    act(() => oldEdit((tx) => tx.setCode("c1", "first draft")));
+    rerender(tree(second));
+    act(() => first.controller.abort());
+    expect(first.detached).toHaveBeenCalledOnce();
+    const latest: {
+      code: string | null | undefined;
+      edit: ReturnType<typeof useDocumentEdit>;
+    } = seen.mock.calls.at(-1)![0];
+    expect(latest.code).toBe("second source");
+    expect(latest.edit).not.toBe(oldEdit);
+    act(() => latest.edit((tx) => tx.setCode("c1", "second draft")));
+    expect(seen.mock.calls.at(-1)![0].code).toBe("second draft");
+    const stale = vi.fn();
+    oldEdit(stale);
+    expect(stale).not.toHaveBeenCalled();
+    seen.mockClear();
+    act(() =>
+      first.store.set(
+        first.source,
+        withCells(document(), {
+          c1: cell("c1", "stale host"),
+        }),
+      ),
+    );
+    expect(seen).not.toHaveBeenCalled();
+  });
+
+  it("requires a document provider", () => {
+    expect(() => renderHook(useDocumentEdit)).toThrow("DocumentProvider");
+  });
+});

package/src/core/document/codemirror.ts

@@ -0,0 +1,311 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * A CodeMirror binding between an editor and one cell's code.
+ *
+ * The editor and the document model each hold the code, and the binding keeps
+ * them equal in both directions.
+ *
+ * - Push. An update that could change the code writes it to the model before
+ *   the update returns, so the model keeps pace with typing.
+ * - Pull. A model change reaches the editor in a microtask, because a view
+ *   cannot dispatch while it is constructing or updating. The binding keeps
+ *   the code both sides last agreed on, so the model's echo of a push is
+ *   skipped while a different value, even one written during that push, is
+ *   pulled. A real change is marked remote, stays out of undo history, and
+ *   waits for IME composition to end.
+ *
+ * The editor shows code through marimo's language adapters, so a markdown
+ * cell edits the text inside `mo.md`. Both directions compare the cell's
+ * Python code rather than the editor text, so an adapter that normalizes what
+ * it reads back leaves typing alone.
+ *
+ * A pulled change replaces only what differs, using `diff` from
+ * `@codemirror/merge`. Separate edits stay separate, so the selection, folds,
+ * and diagnostics between them keep their place. The diff uses a scan limit of
+ * 500, as the merge view does. Past it, the search trades precision for speed,
+ * so a rewrite of a large cell costs microseconds and maps the selection
+ * coarsely instead of blocking the editor.
+ */
+
+import { diff } from "@codemirror/merge";
+import {
+  type ChangeSpec,
+  type EditorState,
+  type Extension,
+  type StateEffect,
+  Transaction,
+} from "@codemirror/state";
+import { type EditorView, ViewPlugin, type ViewUpdate } from "@codemirror/view";
+import type { CellId } from "@marimo-team/client/protocol";
+import { dequal } from "dequal";
+import type { Atom, createStore } from "jotai/vanilla";
+import {
+  languageAdapterFromCode,
+  languageAdapterState,
+  switchLanguage,
+} from "../codemirror/language/extension";
+import { initializeSQLDialect } from "../codemirror/language/languages/sql/sql";
+import {
+  languageMetadataField,
+  setLanguageMetadata,
+} from "../codemirror/language/metadata";
+import type { DocumentModel } from "./model";
+import { selectCellField } from "./select";
+
+/** The diff scan limit, matching the default of the merge view. */
+const SCAN_LIMIT = 500;
+
+interface Binding {
+  readonly store: ReturnType<typeof createStore>;
+  readonly model: DocumentModel;
+  readonly cellId: CellId;
+  /** Undefined while the cell is absent, null while its source is withheld. */
+  readonly code: Atom<string | null | undefined>;
+}
+
+/**
+ * Bind an editor to one cell's code.
+ *
+ * Editors without marimo's language configuration bind their text as is.
+ */
+export function cellCodeBinding({
+  store,
+  model,
+  cellId,
+}: {
+  store: ReturnType<typeof createStore>;
+  model: DocumentModel;
+  cellId: CellId;
+}): Extension {
+  const binding: Binding = {
+    store,
+    model,
+    cellId,
+    code: selectCellField(model.documentAtom, cellId, "code"),
+  };
+  return ViewPlugin.define((view) => new CellCodeSync(view, binding), {
+    eventObservers: {
+      compositionend() {
+        this.compositionEnded();
+      },
+    },
+  });
+}
+
+class CellCodeSync {
+  #scheduled = false;
+  /** A pull arrived during IME composition and runs once it ends. */
+  #deferred = false;
+  /**
+   * A pull is dispatching to the editor.
+   *
+   * A language switch dispatches transactions of its own before the text, and
+   * their in-between states represent the wrong code. Until the switch can be
+   * part of the pull's one remote transaction, this keeps them from pushing.
+   */
+  #pulling = false;
+  /**
+   * The code the editor and the model last agreed on.
+   *
+   * A change the model cannot take clears it, so once the cell is available
+   * again its code is pulled even if it matches what was agreed before.
+   */
+  #synced: string | undefined;
+  #destroyed = false;
+  readonly #view: EditorView;
+  readonly #binding: Binding;
+  readonly #unsubscribe: () => void;
+
+  public constructor(view: EditorView, binding: Binding) {
+    this.#view = view;
+    this.#binding = binding;
+    // A change to the agreed code is the echo of a push and needs no pull.
+    this.#unsubscribe = binding.store.sub(binding.code, () => {
+      if (binding.store.get(binding.code) !== this.#synced) {
+        this.#schedule();
+      }
+    });
+    this.#schedule();
+  }
+
+  public update(update: ViewUpdate) {
+    if (!this.#pulling && changesCode(update)) {
+      this.#push(update.state);
+    }
+    if (this.#deferred && !update.view.composing) {
+      this.#schedule();
+    }
+  }
+
+  /**
+   * Run a pull that waited for the composition.
+   *
+   * A composition that ends without changing the text need not produce an
+   * update. CodeMirror finishes handling the end in a microtask, so the pull
+   * waits for the next task.
+   */
+  public compositionEnded() {
+    if (this.#deferred) {
+      setTimeout(() => this.#schedule(), 0);
+    }
+  }
+
+  public destroy() {
+    this.#destroyed = true;
+    this.#unsubscribe();
+  }
+
+  /** Write the editor's code to the model unless the editor already shows it. */
+  #push(state: EditorState) {
+    const { store, model, cellId, code } = this.#binding;
+    const current = store.get(code);
+    if (typeof current !== "string") {
+      this.#synced = undefined;
+      return;
+    }
+    if (shows(state, current)) {
+      this.#synced = current;
+      return;
+    }
+    const next = codeOf(state);
+    this.#synced = next;
+    store.set(model.editAtom, (tx) => tx.setCode(cellId, next));
+  }
+
+  /** Show the model's code unless the editor already shows it. */
+  #pull() {
+    this.#scheduled = false;
+    const next = this.#binding.store.get(this.#binding.code);
+    if (this.#destroyed || typeof next !== "string") {
+      return;
+    }
+    if (this.#view.composing) {
+      this.#deferred = true;
+      return;
+    }
+    this.#deferred = false;
+    if (next !== this.#synced && !shows(this.#view.state, next)) {
+      this.#pulling = true;
+      try {
+        replaceCode(this.#view, next);
+      } finally {
+        this.#pulling = false;
+      }
+    }
+    this.#synced = next;
+  }
+
+  /** Pull once after the current update, however many changes arrive. */
+  #schedule() {
+    if (!this.#scheduled) {
+      this.#scheduled = true;
+      queueMicrotask(() => this.#pull());
+    }
+  }
+}
+
+/**
+ * The cell code an editor state represents.
+ *
+ * With a language adapter, this is the adapter's Python for the editor text.
+ */
+function codeOf(state: EditorState): string {
+  const text = state.doc.toString();
+  const adapter = state.field(languageAdapterState, false);
+  return adapter
+    ? adapter.transformOut(
+        text,
+        state.field(languageMetadataField, false) ?? {},
+      )[0]
+    : text;
+}
+
+/**
+ * Whether an editor state shows the given cell code.
+ *
+ * An adapter can spell the same code more than one way, such as `mo.md("hi")`
+ * and its canonical form, so a state shows code that reads into its text and
+ * metadata as well as the code its own spelling produces.
+ */
+function shows(state: EditorState, code: string): boolean {
+  if (codeOf(state) === code) {
+    return true;
+  }
+  const adapter = state.field(languageAdapterState, false);
+  if (!adapter) {
+    return false;
+  }
+  const [text, , metadata] = adapter.transformIn(code);
+  return (
+    text === state.doc.toString() &&
+    dequal(metadata, state.field(languageMetadataField, false))
+  );
+}
+
+/**
+ * Whether an update could change the code an editor represents.
+ *
+ * Language metadata, such as an SQL cell's output name, changes the code
+ * without changing the text.
+ */
+function changesCode(update: ViewUpdate): boolean {
+  return (
+    update.docChanged ||
+    update.state.field(languageAdapterState, false) !==
+      update.startState.field(languageAdapterState, false) ||
+    update.state.field(languageMetadataField, false) !==
+      update.startState.field(languageMetadataField, false)
+  );
+}
+
+/**
+ * Show cell code in the language it implies.
+ *
+ * The language switch happens first, in its own transactions, and clears undo
+ * history because its entries edit the previous language's text. The text and
+ * metadata follow in one remote transaction that undo history leaves out. An
+ * SQL dialect depends on the engine in the metadata, so it is set again after.
+ */
+function replaceCode(view: EditorView, code: string) {
+  let text = code;
+  const effects: StateEffect<unknown>[] = [];
+  let sql = false;
+  const adapter = view.state.field(languageAdapterState, false);
+  if (adapter) {
+    const target = languageAdapterFromCode(code);
+    sql = target.type === "sql";
+    if (target.type !== adapter.type) {
+      switchLanguage(view, { language: target.type, keepCodeAsIs: true });
+    }
+    const [inner, , metadata] = target.transformIn(code);
+    text = inner;
+    if (!dequal(metadata, view.state.field(languageMetadataField, false))) {
+      effects.push(setLanguageMetadata.of(metadata));
+    }
+  }
+  const current = view.state.doc.toString();
+  if (current === text && effects.length === 0) {
+    return;
+  }
+  view.dispatch({
+    changes: changesBetween(current, text),
+    effects,
+    annotations: [
+      Transaction.remote.of(true),
+      Transaction.addToHistory.of(false),
+    ],
+  });
+  if (sql && effects.length > 0) {
+    initializeSQLDialect(view);
+  }
+}
+
+/** The changes that turn one text into another, one per edited range. */
+function changesBetween(from: string, to: string): ChangeSpec {
+  return diff(from, to, { scanLimit: SCAN_LIMIT }).map((change) => ({
+    from: change.fromA,
+    to: change.toA,
+    insert: to.slice(change.fromB, change.toB),
+  }));
+}