@marimo-team/islands 0.25.3-dev30 → 0.25.3-dev301

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,158 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type {
+  DocumentModel,
+  DocumentSyncStatus,
+  EditableCell,
+  EditableDocument,
+} from "@/core/notebook/document";
+import { store } from "@/core/state/jotai";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { exportedForTesting, type NotebookState } from "../cells";
+import { applyDocument, flushDocument, followDocument } from "../document";
+import { type CellId, SCRATCH_CELL_ID } from "../ids";
+
+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 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: "_" });
+    expect(state.cellData[b]).toMatchObject({ code: "", name: "_" });
+    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("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 },
+    });
+  });
+
+  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("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>>(),
+  } 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();
+  });
+});

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

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

@@ -0,0 +1,176 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type { CellColumnId } from "@/core/cells/layout";
+import type { CellId } from "../ids";
+import {
+  collapsedAtom,
+  columnOf,
+  deriveLayout,
+  foldSections,
+  foldsAfter,
+  foldsHiding,
+  isColumnId,
+  layoutAtom,
+  newColumnId,
+  visibleNeighbor,
+} from "../layout";
+import type { Outline } from "../outline";
+import { setupNotebook } from "./harness";
+
+const [h1, a, h2, b, c, h1b, d] = ["h1", "a", "h2", "b", "c", "h1b", "d"].map(
+  cellId,
+);
+const ids = [h1, a, h2, b, c, h1b, d];
+
+function heading(level: number): Outline {
+  return { items: [{ name: `H${level}`, level, by: { id: `h${level}` } }] };
+}
+const outlines = new Map<CellId, Outline>([
+  [h1, heading(1)],
+  [h2, heading(2)],
+  [h1b, heading(1)],
+]);
+const outlineOf = (id: CellId) => outlines.get(id) ?? null;
+const column = (id: string, cells: CellId[]) => ({
+  id: id as CellColumnId,
+  ids: cells,
+});
+
+describe("foldSections", () => {
+  it("folds a header until the next header of its level or higher", () => {
+    const folds = foldSections(ids, ids.map(outlineOf), new Set([h2]));
+    expect([...folds]).toEqual([[h2, [b, c]]]);
+  });
+
+  it("hides a folded section inside a folded section", () => {
+    const folds = foldSections(ids, ids.map(outlineOf), new Set([h1, h2]));
+    expect(folds.get(h1)).toEqual([a, h2, b, c]);
+    expect(folds.get(h2)).toEqual([b, c]);
+  });
+
+  it("folds the last section to the end", () => {
+    const folds = foldSections(ids, ids.map(outlineOf), new Set([h1b]));
+    expect(folds.get(h1b)).toEqual([d]);
+  });
+
+  it("ignores cells without a heading", () => {
+    const folds = foldSections([h1, a], [heading(1), null], new Set([a]));
+    expect(folds.size).toBe(0);
+  });
+
+  it("keeps an empty fold for a header with nothing under it yet", () => {
+    const folds = foldSections(
+      [h1, a, h1b],
+      [heading(1), null, heading(1)],
+      new Set([h1b]),
+    );
+    expect(folds.get(h1b)).toEqual([]);
+  });
+});
+
+describe("column ids", () => {
+  it("are told apart from cell ids, as drag and drop needs", () => {
+    expect(isColumnId(newColumnId())).toBe(true);
+    expect(isColumnId(cellId("Hbol"))).toBe(false);
+  });
+});
+
+describe("deriveLayout", () => {
+  const layout = deriveLayout(
+    [column("left", ids), column("right", [])],
+    outlineOf,
+    new Set([h2]),
+  );
+
+  it("lists visible cells per column and keeps empty columns", () => {
+    expect(layout.columns.map((c) => c.id)).toEqual(["left", "right"]);
+    expect(layout.columns[0].visible).toEqual([h1, a, h2, h1b, d]);
+    expect(layout.columns[0].ids).toEqual(ids);
+    expect(layout.columns[1].visible).toEqual([]);
+  });
+
+  it("finds a cell's column, the folds hiding it, and its visible neighbours", () => {
+    expect(columnOf(layout, b).id).toBe("left");
+    expect(() => columnOf(layout, cellId("missing"))).toThrow(
+      "not in the layout",
+    );
+    expect(foldsHiding(layout, c)).toEqual([h2]);
+    expect(foldsHiding(layout, a)).toEqual([]);
+    expect(visibleNeighbor(layout, h2, 1)).toBe(h1b);
+    expect(visibleNeighbor(layout, h1b, -1)).toBe(h2);
+    expect(visibleNeighbor(layout, h1, -1)).toBeUndefined();
+    expect(visibleNeighbor(layout, b, 1)).toBeUndefined();
+  });
+
+  it("names the folds a cell placed after another would fall into", () => {
+    // After the folded header or the last cell of its fold, inside it.
+    expect(foldsAfter(layout, h2)).toEqual([h2]);
+    expect(foldsAfter(layout, c)).toEqual([h2]);
+    // After a visible cell no fold reaches, nothing.
+    expect(foldsAfter(layout, a)).toEqual([]);
+    expect(foldsAfter(layout, h1b)).toEqual([]);
+
+    const nested = deriveLayout(
+      [column("left", ids)],
+      outlineOf,
+      new Set([h1, h2]),
+    );
+    expect(foldsAfter(nested, c)).toEqual([h1, h2]);
+  });
+});
+
+describe("layoutAtom", () => {
+  const stops: Array<() => void> = [];
+  afterEach(() => {
+    for (const stop of stops.splice(0)) {
+      stop();
+    }
+  });
+  const [c1, c2] = ["c1", "c2"].map(cellId);
+
+  /** Deliver `html` as a cell's output, as the kernel's message does. */
+  function output(
+    dispatch: ReturnType<typeof setupNotebook>["dispatch"],
+    id: CellId,
+    html: string,
+  ) {
+    dispatch({
+      type: "handleCellMessage",
+      payload: {
+        cell_id: id,
+        output: {
+          channel: "output",
+          mimetype: "text/html",
+          data: html,
+          timestamp: 1,
+        },
+        console: null,
+        status: "idle",
+        timestamp: 1,
+      },
+    });
+  }
+
+  it("is the same layout after output that changes no heading", () => {
+    const { store, dispatch } = setupNotebook(stops);
+    output(dispatch, c1, "<h1>Title</h1>");
+    const layout = store.get(layoutAtom);
+
+    output(dispatch, c2, "<p>1</p>");
+    output(dispatch, c1, "<h1>Title</h1>");
+
+    expect(store.get(layoutAtom)).toBe(layout);
+  });
+
+  it("follows a heading arriving in a cell's output", () => {
+    const { store, dispatch } = setupNotebook(stops);
+    store.set(collapsedAtom, new Set([c1]));
+    expect(store.get(layoutAtom).columns[0].folds.size).toBe(0);
+
+    output(dispatch, c1, "<h1>Title</h1>");
+
+    expect(store.get(layoutAtom).columns[0].folds.get(c1)).toEqual([c2]);
+  });
+});