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

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

@@ -0,0 +1,410 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { atom } from "jotai";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type {
+  DocumentEdit,
+  DocumentEditor,
+  DocumentModel,
+  DocumentSyncStatus,
+  EditableCell,
+  EditableDocument,
+} from "@/core/notebook/document";
+import {
+  adaptiveLanguageConfiguration,
+  languageAdapterState,
+} from "@/core/codemirror/language/extension";
+import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
+import { store } from "@/core/state/jotai";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { MultiColumn } from "@/utils/id-tree";
+import { isRtcEnabled } from "../../rtc/state";
+import {
+  createNotebookActions,
+  exportedForTesting,
+  type NotebookState,
+  notebookReducer,
+} from "../cells";
+import { applyDocument, flushDocument, withDocumentSave } from "../document";
+import { type CellId, SCRATCH_CELL_ID } from "../ids";
+
+vi.mock("../../rtc/state", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("../../rtc/state")>()),
+  isRtcEnabled: vi.fn(() => false),
+}));
+
+const { initialNotebookState } = exportedForTesting;
+
+const a = cellId("a");
+const b = cellId("b");
+const c = cellId("c");
+const d = cellId("d");
+
+function cell(id: CellId, patch: Partial<EditableCell> = {}): EditableCell {
+  return {
+    id,
+    code: `${id} = 1`,
+    name: "_",
+    config: {
+      column: null,
+      disabled: false,
+      hide_code: false,
+      expand_output: false,
+    },
+    ...patch,
+  };
+}
+
+function documentOf(cells: EditableCell[]): EditableDocument {
+  return {
+    order: cells.map((cell) => cell.id),
+    cells: new Map(cells.map((cell) => [cell.id, cell])),
+  };
+}
+
+/** A mounted cell editor, as the notebook's cells have. */
+function createEditor(content: string) {
+  const state = EditorState.create({
+    doc: content,
+    extensions: [
+      python(),
+      adaptiveLanguageConfiguration({
+        cellId: a,
+        completionConfig: {
+          activate_on_typing: true,
+          signature_hint_on_typing: false,
+          copilot: false,
+          codeium_api_key: null,
+        },
+        hotkeys: new OverridingHotkeyProvider({}),
+        placeholderType: "marimo-import",
+        lspConfig: {},
+      }),
+    ],
+  });
+  return new EditorView({ state, parent: globalThis.document.body });
+}
+
+/** A notebook that follows `cells`, as the session builds one. */
+function notebookOf(cells: EditableCell[]): NotebookState {
+  return applyDocument(initialNotebookState(), documentOf(cells));
+}
+
+describe("applyDocument", () => {
+  it("adds the document's cells to an empty notebook", () => {
+    const state = notebookOf([cell(a), cell(b, { code: "" })]);
+
+    expect(state.cellIds.inOrderIds).toEqual([a, b]);
+    expect(state.cellData[a]).toMatchObject({
+      code: "a = 1",
+      name: "_",
+      edited: false,
+      lastCodeRun: null,
+    });
+    expect(state.cellRuntime[a]).toBeDefined();
+    expect(state.cellHandles[a]).toBeDefined();
+    expect(state.cellData[SCRATCH_CELL_ID]).toBeDefined();
+  });
+
+  it("returns the same state when the document matches", () => {
+    const cells = [cell(a), cell(b)];
+    const state = notebookOf(cells);
+
+    expect(applyDocument(state, documentOf(cells))).toBe(state);
+  });
+
+  it("brings a code change into the cell's editor", () => {
+    const state = notebookOf([cell(a), cell(b)]);
+    const view = createEditor("b = 1");
+    state.cellHandles[b].current = { editorView: view, editorViewOrNull: view };
+    state.cellData[b] = { ...state.cellData[b], lastCodeRun: "b = 1" };
+
+    const next = applyDocument(
+      state,
+      documentOf([cell(a), cell(b, { code: "b = 2" })]),
+    );
+
+    expect(next.cellData[b]).toMatchObject({ code: "b = 2", edited: true });
+    expect(next.cellData[a]).toBe(state.cellData[a]);
+    expect(view.state.doc.toString()).toBe("b = 2");
+  });
+
+  it("leaves an editor that already shows the document's code alone", () => {
+    const state = notebookOf([cell(a)]);
+    // The editor runs ahead of the notebook while this client types.
+    const view = createEditor('mo.md("hi")');
+    state.cellHandles[a].current = { editorView: view, editorViewOrNull: view };
+
+    const next = applyDocument(
+      state,
+      documentOf([cell(a, { code: 'mo.md("hi")' })]),
+    );
+
+    expect(next.cellData[a].code).toBe('mo.md("hi")');
+    expect(view.state.doc.toString()).toBe('mo.md("hi")');
+    expect(view.state.field(languageAdapterState).type).toBe("python");
+  });
+
+  it("leaves an RTC editor to Loro", () => {
+    vi.mocked(isRtcEnabled).mockReturnValueOnce(true);
+    const state = notebookOf([cell(a)]);
+    const view = createEditor("a = 1");
+    state.cellHandles[a].current = { editorView: view, editorViewOrNull: view };
+
+    const next = applyDocument(state, documentOf([cell(a, { code: "a = 2" })]));
+
+    expect(next.cellData[a].code).toBe("a = 2");
+    expect(view.state.doc.toString()).toBe("a = 1");
+  });
+
+  it("keeps the code this client has when the host withholds source", () => {
+    const state = notebookOf([cell(a)]);
+
+    expect(applyDocument(state, documentOf([cell(a, { code: null })]))).toBe(
+      state,
+    );
+  });
+
+  it("follows a cell's name and config", () => {
+    const state = notebookOf([cell(a)]);
+
+    const next = applyDocument(
+      state,
+      documentOf([
+        cell(a, {
+          name: "renamed",
+          config: { ...cell(a).config, hide_code: true },
+        }),
+      ]),
+    );
+
+    expect(next.cellData[a]).toMatchObject({
+      name: "renamed",
+      config: { hide_code: true },
+      edited: false,
+    });
+  });
+
+  it("removes cells the document no longer has", () => {
+    const state = notebookOf([cell(a), cell(b), cell(c)]);
+
+    const next = applyDocument(state, documentOf([cell(a), cell(c)]));
+
+    expect(next.cellIds.inOrderIds).toEqual([a, c]);
+    expect(next.cellData[b]).toBeUndefined();
+    expect(next.cellRuntime[b]).toBeUndefined();
+    expect(next.cellHandles[b]).toBeUndefined();
+  });
+
+  it("regroups columns from the cells' configs, keeping column ids and collapsed sections", () => {
+    const column1 = { ...cell(d).config, column: 1 };
+    const state = notebookOf([
+      cell(a),
+      cell(b),
+      cell(c),
+      cell(d, { config: column1 }),
+    ]);
+    const [first, second] = state.cellIds.getColumns();
+    const collapsed = {
+      ...state,
+      cellIds: new MultiColumn([first.collapse(a, b), second]),
+    };
+
+    const next = applyDocument(
+      collapsed,
+      documentOf([cell(c), cell(a), cell(b), cell(d, { config: column1 })]),
+    );
+
+    expect(next.cellIds.inOrderIds).toEqual([c, a, b, d]);
+    expect(next.cellIds.getColumnIds()).toEqual(
+      collapsed.cellIds.getColumnIds(),
+    );
+    expect(next.cellIds.atOrThrow(0).isCollapsed(a)).toBe(true);
+    expect(next.cellIds.atOrThrow(1).inOrderIds).toEqual([d]);
+  });
+
+  it("keeps a trailing empty column", () => {
+    const state = notebookOf([cell(a)]);
+    const withEmpty = {
+      ...state,
+      cellIds: state.cellIds.addColumn(state.cellIds.atOrThrow(0).id),
+    };
+
+    const next = applyDocument(withEmpty, documentOf([cell(a, { code: "" })]));
+
+    expect(next.cellIds.colLength).toBe(2);
+    expect(next.cellData[a].code).toBe("");
+  });
+});
+
+describe("documentEditMiddleware", () => {
+  type Call = [string, ...unknown[]];
+  let calls: Call[];
+  let state: NotebookState;
+  const editor: DocumentEditor = {
+    setCode: (id, code) => calls.push(["setCode", id, code]),
+    setName: (id, name) => calls.push(["setName", id, name]),
+    updateConfig: (id, patch) => calls.push(["updateConfig", id, patch]),
+    createCell: (cell, placement) =>
+      calls.push(["createCell", cell, placement]),
+    deleteCell: (id) => calls.push(["deleteCell", id]),
+    setOrder: (order) => calls.push(["setOrder", order]),
+  };
+  const model: DocumentModel = {
+    documentAtom: atom(null),
+    editAtom: atom(null, (_get, _set, edit: DocumentEdit) => edit(editor)),
+    syncStatusAtom: atom<DocumentSyncStatus>({ status: "idle" }),
+    flush: vi.fn(),
+    retry: vi.fn(),
+    hold: () => () => undefined,
+    start: vi.fn(),
+  };
+  const actions = createNotebookActions((action) => {
+    state = notebookReducer(state, action);
+  });
+
+  beforeEach(() => {
+    calls = [];
+    state = notebookOf([cell(a), cell(b), cell(c)]);
+    store.set(activeDocumentModelAtom, model);
+  });
+
+  afterEach(() => {
+    store.set(activeDocumentModelAtom, null);
+  });
+
+  it("creates a cell at the end without setting the order", () => {
+    actions.createNewCell({
+      cellId: c,
+      before: false,
+      code: "d = 1",
+      newCellId: d,
+    });
+
+    expect(calls.map((call) => call[0])).toEqual(["createCell"]);
+    expect(calls[0][1]).toMatchObject({ id: d, code: "d = 1" });
+  });
+
+  it("creates a cell elsewhere and sets the order", () => {
+    actions.createNewCell({ cellId: a, before: true, newCellId: d });
+
+    expect(calls.map((call) => call[0])).toEqual(["createCell", "setOrder"]);
+    expect(calls[1][1]).toEqual([d, a, b, c]);
+  });
+
+  it("deletes a cell", () => {
+    actions.deleteCell({ cellId: b });
+
+    expect(calls).toEqual([["deleteCell", b]]);
+  });
+
+  it("writes code, name, and config", () => {
+    actions.updateCellCode({
+      cellId: a,
+      code: "a = 2",
+      formattingChange: false,
+    });
+    actions.updateCellName({ cellId: a, name: "first" });
+    actions.updateCellConfig({ cellId: a, config: { hide_code: true } });
+
+    expect(calls).toEqual([
+      ["setCode", a, "a = 2"],
+      ["setName", a, "first"],
+      ["updateConfig", a, { hide_code: true }],
+    ]);
+  });
+
+  it("sets the order when a cell moves", () => {
+    actions.moveCell({ cellId: c, before: true });
+
+    expect(calls).toEqual([["setOrder", [a, c, b]]]);
+  });
+
+  it("sets the column of a cell that changes column", () => {
+    state = notebookOf([
+      cell(a),
+      cell(b),
+      cell(c, { config: { ...cell(c).config, column: 1 } }),
+    ]);
+
+    actions.moveCell({ cellId: b, direction: "right" });
+
+    expect(state.cellIds.atOrThrow(1).inOrderIds).toEqual([b, c]);
+    expect(calls).toEqual([["updateConfig", b, { column: 1 }]]);
+  });
+
+  it("clears columns merged into one", () => {
+    state = notebookOf([
+      cell(a),
+      cell(b),
+      cell(c, { config: { ...cell(c).config, column: 1 } }),
+    ]);
+
+    actions.mergeAllColumns();
+
+    expect(calls).toEqual([["updateConfig", c, { column: null }]]);
+  });
+
+  it("ignores the kernel's cells and runtime state", () => {
+    actions.prepareForRun({ cellId: a });
+    actions.setCells([state.cellData[a]]);
+
+    expect(calls).toEqual([]);
+  });
+
+  it("leaves the scratchpad alone", () => {
+    actions.updateCellCode({
+      cellId: SCRATCH_CELL_ID,
+      code: "scratch",
+      formattingChange: false,
+    });
+
+    expect(calls).toEqual([]);
+  });
+});
+
+describe("flushing before host commands", () => {
+  const model = {
+    syncStatusAtom: atom<DocumentSyncStatus>({ status: "idle" }),
+    flush: vi.fn<() => Promise<never>>(),
+    retry: vi.fn<() => Promise<never>>(),
+    hold: () => () => undefined,
+  } as unknown as DocumentModel;
+
+  beforeEach(() => {
+    vi.mocked(model.flush).mockReset();
+    vi.mocked(model.retry).mockReset();
+    store.set(activeDocumentModelAtom, model);
+  });
+
+  afterEach(() => {
+    store.set(activeDocumentModelAtom, null);
+  });
+
+  it("flushes the session's model, or nothing without one", async () => {
+    vi.mocked(model.flush).mockResolvedValue(undefined as never);
+    await flushDocument();
+    expect(model.flush).toHaveBeenCalledOnce();
+
+    store.set(activeDocumentModelAtom, null);
+    await expect(flushDocument()).resolves.toBeUndefined();
+  });
+
+  it("saves after a flush and resumes the model when the flush failed", async () => {
+    vi.mocked(model.flush).mockRejectedValue(new Error("conflict"));
+    vi.mocked(model.retry).mockResolvedValue(undefined as never);
+    const save = vi.fn().mockResolvedValue("saved");
+
+    await expect(withDocumentSave(save)).resolves.toBe("saved");
+    expect(save).toHaveBeenCalledOnce();
+    expect(model.retry).toHaveBeenCalledOnce();
+
+    vi.mocked(model.flush).mockResolvedValue(undefined as never);
+    await withDocumentSave(save);
+    expect(model.retry).toHaveBeenCalledOnce();
+  });
+});

package/src/core/cells/__tests__/flush-document.test.ts

@@ -0,0 +1,40 @@
+/* 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 {
+  type FakeEdit,
+  fakeBackend,
+} from "@/core/notebook/replica/__tests__/fixtures";
+import type { Document } from "@/core/notebook/replica/resources/document";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { store } from "@/core/state/jotai";
+import { flushDocument } from "../document";
+
+const lifetimes: AbortController[] = [];
+
+afterEach(() => {
+  lifetimes.splice(0).forEach((controller) => controller.abort());
+  store.set(activeDocumentModelAtom, null);
+});
+
+describe("flushDocument", () => {
+  it("sends nothing while the document is unavailable", async () => {
+    const edit = vi.fn<FakeEdit>();
+    const model = createDocumentModel({
+      store,
+      documentAtom: atom<Document | null>(null),
+      backendAtom: atom(fakeBackend({ edit })),
+    });
+    const controller = new AbortController();
+    lifetimes.push(controller);
+    model.start({ signal: controller.signal });
+    store.set(activeDocumentModelAtom, model);
+
+    expect(store.get(model.syncStatusAtom)).toEqual({ status: "unavailable" });
+    await expect(model.flush()).rejects.toThrow("unavailable");
+    await expect(flushDocument()).resolves.toBeUndefined();
+    expect(edit).not.toHaveBeenCalled();
+  });
+});

package/src/core/cells/__tests__/notebook-outline.test.ts

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { createStore } from "jotai";
-import { afterEach, describe, expect, it, vi } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { cellId } from "@/__tests__/branded";
 import {
@@ -10,7 +10,6 @@ import {
   notebookOutline,
   notebookReducer,
 } from "../cells";
-import { exportedForTesting } from "../document-changes";
 import type { Outline } from "../outline";
 
 const first = cellId("first");
@@ -40,10 +39,6 @@ function setup() {
   return { store, actions };
 }
 
-afterEach(() => {
-  exportedForTesting.cancelPendingChanges();
-});
-
 describe("notebookOutline", () => {
   it("does not notify outline subscribers when cell code changes", () => {
     const { store, actions } = setup();

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);
+  });
+});