@marimo-team/islands 0.25.3-dev20 → 0.25.3-dev200

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

@@ -2,6 +2,7 @@
 
 import { describe, expect, it } from "vitest";
 import { StatefulOutputMessage } from "@/components/editor/output/ansi-reduce";
+import type { ConsoleOutputMessage } from "@/core/cells/types";
 import type { OutputMessage } from "@/core/kernel/messages";
 import {
   collapseConsoleOutputs,
@@ -279,6 +280,59 @@ describe("collapseConsoleOutputs", () => {
     expect(result[0].data).toContain("Streaming output truncated");
   });
 
+  it("keeps the first id when collapsing outputs", () => {
+    const consoleOutputs = [
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "Hello ",
+        timestamp: 0,
+        id: 7,
+      },
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "World",
+        timestamp: 0,
+        id: 8,
+      },
+    ] satisfies ConsoleOutputMessage[];
+    const result = collapseConsoleOutputs(consoleOutputs);
+    expect(result).toHaveLength(1);
+    expect(result[0].data).toBe("Hello World");
+    expect((result[0] as ConsoleOutputMessage).id).toBe(7);
+  });
+
+  it("keeps ids on outputs that survive truncation", () => {
+    const consoleOutputs = [
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "one\ntwo",
+        timestamp: 0,
+        id: 1,
+      },
+      {
+        mimetype: "text/plain",
+        channel: "stderr",
+        data: "three\nfour",
+        timestamp: 0,
+        id: 2,
+      },
+      {
+        mimetype: "text/plain",
+        channel: "output",
+        data: "five",
+        timestamp: 0,
+        id: 3,
+      },
+    ] satisfies ConsoleOutputMessage[];
+    const result = collapseConsoleOutputs(consoleOutputs, 3);
+    const ids = result.map((output) => (output as ConsoleOutputMessage).id);
+    expect(result[0].data).toContain("Streaming output truncated");
+    expect(ids).toEqual([undefined, 2, 3]);
+  });
+
   describe("ANSI escape sequences", () => {
     it("should handle cursor movement with collapse", () => {
       const consoleOutputs: OutputMessage[] = [

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

@@ -0,0 +1,341 @@
+/* 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 { isRtcEnabled } from "../../rtc/state";
+import {
+  createNotebookActions,
+  exportedForTesting,
+  type NotebookState,
+  notebookReducer,
+} from "../cells";
+import {
+  applyDocument,
+  flushDocument,
+  followDocument,
+  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");
+
+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).toEqual([a, b]);
+    expect(state.cellData[a]).toMatchObject({
+      code: "a = 1",
+      name: "_",
+      edited: true,
+      lastCodeRun: null,
+    });
+    expect(state.cellData[b].edited).toBe(false);
+    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: true,
+    });
+  });
+
+  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).toEqual([a, c]);
+    expect(next.cellData[b]).toBeUndefined();
+    expect(next.cellRuntime[b]).toBeUndefined();
+    expect(next.cellHandles[b]).toBeUndefined();
+  });
+});
+
+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("writes code", () => {
+    actions.updateCellCode({
+      cellId: a,
+      code: "a = 2",
+      formattingChange: false,
+    });
+
+    expect(calls).toEqual([["setCode", a, "a = 2"]]);
+  });
+
+  it("ignores the document and the kernel's run state", () => {
+    actions.prepareForRun({ cellId: a });
+    actions.setLastRun({
+      lastCodeRun: { [a]: "a = 0" },
+      lastExecutionTime: {},
+    });
+    actions.resetCellRuntimes();
+    actions.reconcileDocument(documentOf([cell(a)]));
+
+    expect(calls).toEqual([]);
+  });
+
+  it("leaves the scratchpad alone", () => {
+    actions.updateCellCode({
+      cellId: SCRATCH_CELL_ID,
+      code: "scratch",
+      formattingChange: false,
+    });
+
+    expect(calls).toEqual([]);
+  });
+});
+
+describe("followDocument", () => {
+  it("reconciles the document now and on each change, until aborted", () => {
+    const documentAtom = atom<EditableDocument | null>(null);
+    const reconcile = vi.fn();
+    const controller = new AbortController();
+
+    followDocument({
+      store,
+      model: { documentAtom },
+      reconcile,
+      signal: controller.signal,
+    });
+    expect(reconcile).not.toHaveBeenCalled();
+
+    const first = documentOf([cell(a)]);
+    store.set(documentAtom, first);
+    expect(reconcile).toHaveBeenCalledWith(first);
+
+    controller.abort();
+    store.set(documentAtom, documentOf([cell(b)]));
+    expect(reconcile).toHaveBeenCalledOnce();
+  });
+});
+
+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__/editor.test.ts

@@ -0,0 +1,57 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { createElement, type PropsWithChildren } from "react";
+import { describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import {
+  hideShownCodeAtom,
+  isCodeHidden,
+  restoredEditorStatesAtom,
+  showHiddenCodeAtom,
+  shownHiddenCodeAtom,
+  takeRestoredEditorStateAtom,
+  useRestoredEditorState,
+} from "../editor";
+
+const a = cellId("a");
+const b = cellId("b");
+
+describe("hidden code showing for now", () => {
+  it("hides code unless the cell is showing it", () => {
+    const shown = new Set([a]);
+    expect(isCodeHidden(true, shown, a)).toBe(false);
+    expect(isCodeHidden(true, shown, b)).toBe(true);
+    expect(isCodeHidden(false, shown, b)).toBe(false);
+    expect(isCodeHidden(undefined, shown, b)).toBe(false);
+  });
+
+  it("shows a cell's code and hides it again, saying whether it was showing", () => {
+    const store = createStore();
+    store.set(showHiddenCodeAtom, { cellId: a });
+    expect(store.get(shownHiddenCodeAtom).has(a)).toBe(true);
+    expect(store.set(hideShownCodeAtom, { cellId: a })).toBe(true);
+    expect(store.get(shownHiddenCodeAtom).has(a)).toBe(false);
+    expect(store.set(hideShownCodeAtom, { cellId: a })).toBe(false);
+  });
+});
+
+describe("restored editor states", () => {
+  it("hands a state over once", () => {
+    const store = createStore();
+    store.set(restoredEditorStatesAtom, new Map([[a, { doc: "x = 1" }]]));
+    const wrapper = ({ children }: PropsWithChildren) =>
+      createElement(Provider, { store }, children);
+    const { result } = renderHook(() => useRestoredEditorState(a), {
+      wrapper,
+    });
+    expect(result.current).toEqual({ doc: "x = 1" });
+
+    act(() => store.set(takeRestoredEditorStateAtom, { cellId: a }));
+
+    expect(result.current).toBeNull();
+    expect(store.get(restoredEditorStatesAtom).size).toBe(0);
+  });
+});

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__/harness.ts

@@ -0,0 +1,106 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A store whose notebook follows a document model, as the session wires it. */
+
+import { history } from "@codemirror/commands";
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { createStore } from "jotai";
+import { mockDocumentModel } from "@/__mocks__/document";
+import type { CellId } from "@/core/cells/ids";
+import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
+import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
+import {
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
+import type { Document } from "@/core/notebook/replica";
+import { exportedForTesting, notebookAtom, notebookReducer } from "../cells";
+import { followDocument } from "../document";
+import { layoutAtom } from "../layout";
+
+export function setupNotebook(
+  stops: Array<() => void>,
+  cells?: Record<string, Protocol.Cell | null>,
+) {
+  const store = createStore();
+  store.set(notebookAtom, exportedForTesting.initialNotebookState());
+  const { model, documentAtom, stop } = mockDocumentModel(
+    store,
+    cells ? withCells(documentOf(), cells) : documentOf(),
+  );
+  const controller = new AbortController();
+  stops.push(stop, () => controller.abort());
+  followDocument({
+    store,
+    model,
+    reconcile: (document) =>
+      store.set(notebookAtom, (state) =>
+        notebookReducer(state, {
+          type: "reconcileDocument",
+          payload: document,
+        }),
+      ),
+    signal: controller.signal,
+  });
+  return {
+    store,
+    model,
+    /** Deliver a document from the host. Edits pending here stay on top. */
+    host: (document: Document) => store.set(documentAtom, document),
+    order: () => store.get(notebookAtom).cellIds,
+    /** The cells of each column, as the page lays them out. */
+    columns: () => store.get(layoutAtom).columns.map((column) => column.ids),
+    /** The columns' ids, as the page keys them. */
+    columnIds: () => store.get(layoutAtom).columns.map((column) => column.id),
+    notebook: () => store.get(notebookAtom),
+    cellIn: (id: string) => store.get(model.documentAtom)?.cells.get(id),
+    /** Give a cell an editor, as mounting it would, with a cell's language setup. */
+    mountEditor: (id: string) => {
+      const view = new EditorView({
+        state: EditorState.create({
+          doc: store.get(model.documentAtom)?.cells.get(id)?.code ?? "",
+          extensions: [
+            history(),
+            python(),
+            adaptiveLanguageConfiguration({
+              cellId: id as CellId,
+              completionConfig: {
+                activate_on_typing: true,
+                signature_hint_on_typing: false,
+                copilot: false,
+                codeium_api_key: null,
+              },
+              hotkeys: new OverridingHotkeyProvider({}),
+              placeholderType: "marimo-import",
+              lspConfig: {},
+            }),
+          ],
+        }),
+      });
+      store.get(notebookAtom).cellHandles[id as never].current = {
+        editorView: view,
+        editorViewOrNull: view,
+      };
+      return view;
+    },
+    /** Give a cell's output a heading, as the kernel's outline would. */
+    heading: (id: string, level: number) =>
+      store.set(notebookAtom, (state) => ({
+        ...state,
+        cellRuntime: {
+          ...state.cellRuntime,
+          [id]: {
+            ...state.cellRuntime[id as never],
+            outline: { items: [{ name: `H${level}`, level, by: { id } }] },
+          },
+        },
+      })),
+    visible: () => store.get(layoutAtom).columns[0].visible,
+    /** Apply a reducer action, as `useReducerActions()` does. */
+    dispatch: (action: Parameters<typeof notebookReducer>[1]) =>
+      store.set(notebookAtom, (state) => notebookReducer(state, action)),
+  };
+}