@marimo-team/frontend 0.25.3-dev13 → 0.25.3-dev134

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,
+} from "..";
+import { cell, CONFIG, document, withCells } from "./fixtures";
+import { type FakeEdit, fakeBackend } from "../../host/__tests__/fixtures";
+
+const lifetimes: AbortController[] = [];
+
+afterEach(() => {
+  act(() => lifetimes.splice(0).forEach((controller) => controller.abort()));
+});
+
+function setup(
+  initial: Document | null = document(),
+  edit: FakeEdit = () => 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,
+    documentAtom: source,
+    backendAtom: atom(fakeBackend({ 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<FakeEdit>()
+      .mockResolvedValueOnce({
+        ok: true,
+        status: 200,
+        data: { document: accepted.revision, cells: {} },
+      })
+      .mockImplementation(
+        () => new Deferred<Awaited<ReturnType<FakeEdit>>>().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/__tests__/lifecycle-scenarios.test.ts

@@ -0,0 +1,61 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { vi } from "vitest";
+import { test } from "./scenario-fixture";
+
+test("closing a session before the debounce prevents submission", async ({
+  expect,
+  session: { user, notebook, controller },
+}) => {
+  user.changeCode("c3", "local");
+  const reason = new Error("Editing session ended");
+  controller.abort(reason);
+  await vi.advanceTimersByTimeAsync(1000);
+  expect(notebook.edit).not.toHaveBeenCalled();
+  expect(vi.getTimerCount()).toBe(0);
+  await expect(user.sync()).rejects.toBe(reason);
+});
+
+test("closing a session before a queued flush starts prevents submission", async ({
+  expect,
+  session: { user, notebook, controller },
+}) => {
+  user.changeCode("c3", "local");
+  const flushed = user.sync();
+  const reason = new Error("Editing session ended");
+  controller.abort(reason);
+  await expect(flushed).rejects.toBe(reason);
+  expect(notebook.edit).not.toHaveBeenCalled();
+  expect(vi.getTimerCount()).toBe(0);
+});
+
+test.for(["request", "confirmation", "retry"] as const)(
+  "closing a session during %s stops synchronization",
+  async (phase, { expect, session: { user, host, notebook, controller } }) => {
+    user.changeCode("c3", "local");
+    const flushed = user.sync();
+    const request = await host.nextRequest();
+    if (phase === "confirmation") {
+      request.process();
+      request.deliverResponse();
+    } else if (phase === "retry") {
+      request.loseResponse();
+    }
+    await vi.advanceTimersByTimeAsync(0);
+    const reason = new Error("Editing session ended");
+    controller.abort(reason);
+    await expect(flushed).rejects.toBe(reason);
+    if (phase === "request") {
+      expect(request.signal.aborted).toBe(true);
+    }
+
+    host.changeCell("c1", { code: "later" });
+    host.deliverSnapshot();
+    user.changeCode("c3", "after closing");
+    await vi.advanceTimersByTimeAsync(1000);
+    expect(notebook.edit).toHaveBeenCalledOnce();
+    expect(vi.getTimerCount()).toBe(0);
+    await expect(user.sync()).rejects.toBe(reason);
+    await expect(user.retry()).rejects.toBe(reason);
+  },
+);

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

@@ -0,0 +1,191 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Result } from "@marimo-team/client";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, createStore } from "jotai/vanilla";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import type { NotebookBackend } from "../../host/backend";
+import type { Document } from "../../host/resources/document";
+import { createDocumentModel } from "../model";
+import { type FakeEdit, fakeBackend } from "../../host/__tests__/fixtures";
+import { cell, document, withCells } from "./fixtures";
+
+function response(confirmed: Document): Result<Protocol.EditResponseOutput> {
+  return {
+    ok: true,
+    status: 200,
+    data: { document: confirmed.revision, cells: {} },
+  };
+}
+
+// Observe rejections before advancing timers to avoid unhandled promises.
+async function expectFailure(promise: Promise<unknown>, message: string) {
+  await expect(promise).rejects.toThrow(message);
+}
+
+async function expectAbort(promise: Promise<unknown>, reason: Error) {
+  await expect(promise).rejects.toBe(reason);
+}
+
+/** An edit the host received, with its notebook and options. */
+interface SentEdit {
+  notebookId: Protocol.NotebookId;
+  changes: readonly Protocol.Change[];
+  idempotencyKey?: string;
+  signal?: AbortSignal;
+}
+
+function setup({ connected = true } = {}) {
+  const store = createStore();
+  const source = atom<Document | null>(document());
+  const unmounted = vi.fn();
+  source.onMount = () => unmounted;
+  const replies: Deferred<Result<Protocol.EditResponseOutput>>[] = [];
+  const sent = vi.fn((_edit: SentEdit) => {
+    const reply = new Deferred<Result<Protocol.EditResponseOutput>>();
+    replies.push(reply);
+    return reply.promise;
+  });
+  const backend = fakeBackend({
+    edit: (notebookId, changes, options) =>
+      sent({ notebookId, changes, ...options }),
+  });
+  const backendAtom = atom<NotebookBackend | null>(connected ? backend : null);
+  const controller = new AbortController();
+  lifetimes.push(controller);
+  const model = createDocumentModel({
+    store,
+    documentAtom: source,
+    backendAtom,
+    signal: controller.signal,
+  });
+  return {
+    store,
+    source,
+    backend,
+    backendAtom,
+    model,
+    sent,
+    replies,
+    controller,
+    unmounted,
+    edit: (code: string) =>
+      store.set(model.editAtom, (tx) => tx.setCode("c1", code)),
+    status: () => store.get(model.syncStatusAtom),
+    value: () => store.get(model.documentAtom)?.cells.get("c1"),
+  };
+}
+
+const lifetimes: AbortController[] = [];
+beforeEach(() => vi.useFakeTimers());
+afterEach(() => {
+  lifetimes.splice(0).forEach((controller) => controller.abort());
+  vi.useRealTimers();
+});
+
+describe("document synchronization", () => {
+  it("rejects an aborted lifetime before subscribing or sending", () => {
+    const store = createStore();
+    const source = atom<Document | null>(document());
+    const mounted = vi.fn();
+    source.onMount = mounted;
+    const controller = new AbortController();
+    const reason = new Error("Session ended");
+    controller.abort(reason);
+    const edit = vi.fn<FakeEdit>();
+    expect(() =>
+      createDocumentModel({
+        store,
+        documentAtom: source,
+        backendAtom: atom(fakeBackend({ edit })),
+        signal: controller.signal,
+      }),
+    ).toThrow(reason);
+    expect(mounted).not.toHaveBeenCalled();
+    expect(edit).not.toHaveBeenCalled();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("fails without repeating a request when the response names another notebook", async () => {
+    const { sent, replies, edit, model, status } = setup();
+    edit("local");
+    const rejected = expectFailure(model.flush(), "another notebook");
+    await vi.advanceTimersByTimeAsync(0);
+    replies[0].resolve(response(document("n2")));
+    await rejected;
+    await vi.advanceTimersByTimeAsync(5000);
+    expect(sent).toHaveBeenCalledOnce();
+    expect(status()).toMatchObject({ status: "error" });
+  });
+
+  it("ignores a late response and source update after its lifetime ends", async () => {
+    const {
+      store,
+      source,
+      model,
+      replies,
+      sent,
+      edit,
+      controller,
+      unmounted,
+      value,
+    } = setup();
+    edit("local");
+    const reason = new Error("Editing session ended");
+    const rejected = expectAbort(model.flush(), reason);
+    await vi.advanceTimersByTimeAsync(0);
+    controller.abort(reason);
+    await rejected;
+    const accepted = withCells(document(), { c1: cell("c1", "local") });
+    replies[0].resolve(response(accepted));
+    store.set(source, accepted);
+    const staleEdit = vi.fn();
+    store.set(model.editAtom, staleEdit);
+    await vi.advanceTimersByTimeAsync(1000);
+    expect(unmounted).toHaveBeenCalledOnce();
+    expect(staleEdit).not.toHaveBeenCalled();
+    expect(value()).toBeUndefined();
+    expect(sent).toHaveBeenCalledOnce();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("aborts an in-flight edit when notebook identity changes", async () => {
+    const { store, source, sent, edit, model, status } = setup();
+    edit("local");
+    const flushed = model.flush();
+    const rejected = expectFailure(flushed, "notebook changed");
+    await vi.advanceTimersByTimeAsync(0);
+    store.set(source, document("n2"));
+    await rejected;
+    expect(sent.mock.calls[0][0].signal?.aborted).toBe(true);
+    expect(status()).toMatchObject({ status: "error", retryable: false });
+  });
+
+  it("aborts an in-flight edit when the backend changes and never sends it to the new one", async () => {
+    const { store, backendAtom, sent, edit, model } = setup();
+    edit("local");
+    const rejected = expectFailure(model.flush(), "notebook changed");
+    await vi.advanceTimersByTimeAsync(0);
+    const other = vi.fn<FakeEdit>();
+    store.set(backendAtom, fakeBackend({ edit: other }));
+    await rejected;
+    expect(sent).toHaveBeenCalledOnce();
+    expect(sent.mock.calls[0][0].signal?.aborted).toBe(true);
+    expect(other).not.toHaveBeenCalled();
+  });
+
+  it("keeps edits as drafts until a backend is connected", async () => {
+    const { store, backend, backendAtom, sent, edit, status, value } = setup({
+      connected: false,
+    });
+    edit("local");
+    await vi.advanceTimersByTimeAsync(1000);
+    expect(status()).toEqual({ status: "unavailable" });
+    expect(value()?.code).toBe("local");
+
+    store.set(backendAtom, backend);
+    await vi.advanceTimersByTimeAsync(400);
+    expect(sent).toHaveBeenCalledOnce();
+  });
+});

package/src/core/document/__tests__/protocol-fixtures.ts

@@ -0,0 +1,38 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { readFileSync } from "node:fs";
+import NodeUrl from "node:url";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { invariant } from "@/utils/invariant";
+import { fixtureEvents } from "../../host/__tests__/fixtures";
+import type { Document } from "../../host/resources/document";
+
+export function protocolFixture(name: string): unknown {
+  return JSON.parse(
+    readFileSync(
+      new NodeUrl.URL(
+        `../../../../../protocols/host/fixtures/${name}`,
+        import.meta.url,
+      ),
+      "utf8",
+    ),
+  );
+}
+
+export function initialDocument(): Document {
+  const cells = new Map<Protocol.CellId, Protocol.Cell>();
+  let revision: Protocol.Document | undefined;
+  for (const event of fixtureEvents("notebook-events.sse")) {
+    if (event.event === "ready") {
+      break;
+    }
+    if (event.event === "cell") {
+      cells.set(event.data.cell.id, event.data.cell);
+    }
+    if (event.event === "document") {
+      revision = event.data.document;
+    }
+  }
+  invariant(revision !== undefined, "Fixture has no document");
+  return { revision, cells };
+}