@marimo-team/islands 0.25.3-dev23 → 0.25.3-dev231

package/src/core/notebook/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 "../../replica/__tests__/fixtures";
+import type { Document } from "../../replica/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 };
+}

package/src/core/notebook/document/__tests__/recovery-scenarios.test.ts

@@ -0,0 +1,230 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { vi } from "vitest";
+import { test } from "./scenario-fixture";
+
+test.for(["response", "document"] as const)(
+  "keeps a conflicting batch intact and retries explicitly with %s first",
+  async (first, { expect, session: { user, host, notebook } }) => {
+    // The first change is valid; the second conflicts. Neither may commit.
+    user.changeCells({
+      c1: { code: "other local edit" },
+      c3: { code: "local" },
+    });
+    const synced = user.sync();
+    const request = await host.nextRequest();
+    host.changeCell("c3", { code: "remote", name: "remote name" });
+    const remote = host.document();
+    expect(user.cell("c3")?.code).toBe("local");
+    request.process();
+    expect(host.document()).toEqual(remote);
+    user.changeCode("c3", "updated local");
+
+    if (first === "document") {
+      host.deliverDocument();
+    }
+    request.deliverResponse();
+    await expect(synced).rejects.toThrow("conflict");
+    expect(user.status()).toMatchObject({ status: "error", retryable: true });
+    if (first === "response") {
+      host.deliverDocument();
+    }
+    await vi.advanceTimersByTimeAsync(2000);
+    expect(notebook.edit).toHaveBeenCalledOnce();
+    expect(user.cell("c3")?.code).toBe("updated local");
+    expect(user.cell("c1")?.code).toBe("other local edit");
+
+    const retried = user.retry();
+    const retry = await host.nextRequest();
+    expect(retry.idempotencyKey).not.toBe(request.idempotencyKey);
+    retry.process();
+    retry.deliverResponse();
+    host.deliverDocument();
+    await expect(retried).resolves.toEqual(host.document().revision);
+    expect(host.document().cells.get("c1")?.code).toBe("other local edit");
+    expect(host.document().cells.get("c3")).toMatchObject({
+      code: "updated local",
+      name: "remote name",
+    });
+    expect(user.cell("c3")).toMatchObject({
+      code: "updated local",
+      name: "remote name",
+    });
+    expect(user.status()).toEqual({ status: "idle" });
+  },
+);
+
+test("waits for the conflicting revision before allowing a retry", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "local");
+  const synced = user.sync();
+  const request = await host.nextRequest();
+  host.changeCell("c3", { code: "remote" });
+  request.process();
+  request.deliverResponse();
+  await expect(synced).rejects.toThrow("conflict");
+  await expect(user.retry()).rejects.toThrow("Wait for the host");
+  expect(notebook.edit).toHaveBeenCalledOnce();
+  expect(user.cell("c3")?.code).toBe("local");
+
+  host.deliverDocument();
+  const retried = user.retry();
+  const retry = await host.nextRequest();
+  retry.process();
+  retry.deliverResponse();
+  host.deliverDocument();
+  await expect(retried).resolves.toEqual(host.document().revision);
+  expect(host.document().cells.get("c3")?.code).toBe("local");
+  expect(user.status()).toEqual({ status: "idle" });
+});
+
+test("replays a lost response without committing twice or losing newer typing", async ({
+  expect,
+  session: { user, host },
+}) => {
+  user.changeCode("c3", "sent");
+  const flushed = user.sync();
+  const request = await host.nextRequest();
+  request.process();
+  const committed = host.document();
+  host.deliverDocument();
+  user.changeCode("c3", "newer");
+  request.loseResponse();
+
+  await vi.advanceTimersByTimeAsync(500);
+  const retry = await host.nextRequest();
+  expect(retry.idempotencyKey).toBe(request.idempotencyKey);
+  expect(retry.changes).toEqual(request.changes);
+  retry.process();
+  expect(host.document()).toEqual(committed);
+  retry.deliverResponse();
+
+  const next = await host.nextRequest();
+  expect(next.idempotencyKey).not.toBe(request.idempotencyKey);
+  next.process();
+  next.deliverResponse();
+  host.deliverDocument();
+  await expect(flushed).resolves.toEqual(host.document().revision);
+  expect(host.document().revision.version).toBe(committed.revision.version + 1);
+  expect(host.document().cells.get("c3")?.code).toBe("newer");
+  expect(user.cell("c3")?.code).toBe("newer");
+});
+
+test("ignores a timed-out attempt that completes after its retry", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "local");
+  const flushed = user.sync();
+  const first = await host.nextRequest();
+  await vi.advanceTimersByTimeAsync(10_500);
+  const retry = await host.nextRequest();
+  expect(first.signal.aborted).toBe(true);
+  expect(retry.idempotencyKey).toBe(first.idempotencyKey);
+  retry.process();
+  const committed = host.document();
+  retry.deliverResponse();
+  host.deliverDocument();
+  await expect(flushed).resolves.toEqual(committed.revision);
+  first.process();
+  first.deliverResponse();
+  await vi.advanceTimersByTimeAsync(0);
+  expect(host.document()).toEqual(committed);
+  expect(user.cell("c3")?.code).toBe("local");
+  expect(notebook.edit).toHaveBeenCalledTimes(2);
+});
+
+test("a manual retry reuses the committed edit after automatic retries give up", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "sent");
+  const synced = user.sync();
+  const request = await host.nextRequest();
+  request.process();
+  const committed = host.document();
+  request.loseResponse();
+  for (const delay of [500, 1000]) {
+    await vi.advanceTimersByTimeAsync(delay);
+    const retry = await host.nextRequest();
+    retry.process();
+    retry.loseResponse();
+  }
+  await expect(synced).rejects.toThrow("Connection lost");
+  expect(user.status()).toMatchObject({ status: "error", retryable: true });
+  expect(notebook.edit).toHaveBeenCalledTimes(3);
+  expect(host.document()).toEqual(committed);
+
+  const retried = user.retry();
+  const retry = await host.nextRequest();
+  expect(retry.idempotencyKey).toBe(request.idempotencyKey);
+  expect(retry.changes).toEqual(request.changes);
+  retry.process();
+  retry.deliverResponse();
+  host.deliverDocument();
+  await expect(retried).resolves.toEqual(committed.revision);
+  expect(host.document()).toEqual(committed);
+  expect(user.status()).toEqual({ status: "idle" });
+  expect(notebook.edit).toHaveBeenCalledTimes(4);
+});
+
+test("refuses a manual retry after the safe retry window while retaining the draft", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "sent");
+  const synced = user.sync();
+  const request = await host.nextRequest();
+  request.process();
+  const committed = host.document();
+  request.loseResponse();
+  for (const delay of [500, 1000]) {
+    await vi.advanceTimersByTimeAsync(delay);
+    const retry = await host.nextRequest();
+    retry.process();
+    retry.loseResponse();
+  }
+  await expect(synced).rejects.toThrow("Connection lost");
+  expect(user.status()).toMatchObject({ status: "error", retryable: true });
+
+  await vi.advanceTimersByTimeAsync(60_000);
+  await expect(user.retry()).rejects.toThrow("retry window expired");
+  expect(user.status()).toMatchObject({ status: "error", retryable: false });
+  expect(notebook.edit).toHaveBeenCalledTimes(3);
+  expect(host.document()).toEqual(committed);
+  expect(user.cell("c3")?.code).toBe("sent");
+});
+
+test("keeps the whole draft when another writer deletes an edited cell", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "keep this");
+  const synced = user.sync();
+  const request = await host.nextRequest();
+  host.deleteCell("c3");
+  request.process();
+  request.deliverResponse();
+  host.deliverDocument();
+  await expect(synced).rejects.toThrow("conflict");
+
+  user.changeCode("c1", "other edit");
+  await expect(user.retry()).rejects.toThrow("unavailable");
+  await vi.advanceTimersByTimeAsync(1000);
+  expect(notebook.edit).toHaveBeenCalledOnce();
+  expect(host.document().cells.get("c1")?.code).toBe("import marimo as mo");
+
+  host.createCell({ id: "c3", code: "restored" });
+  host.deliverDocument();
+  expect(user.cell("c3")?.code).toBe("keep this");
+  const retried = user.retry();
+  const retry = await host.nextRequest();
+  retry.process();
+  retry.deliverResponse();
+  host.deliverDocument();
+  await expect(retried).resolves.toEqual(host.document().revision);
+  expect(host.document().cells.get("c3")?.code).toBe("keep this");
+  expect(host.document().cells.get("c1")?.code).toBe("other edit");
+});

package/src/core/notebook/document/__tests__/scenario-fixture.ts

@@ -0,0 +1,43 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createStore } from "jotai/vanilla";
+import { test as baseTest, vi } from "vitest";
+import type { Document } from "../../replica/resources/document";
+import { createDocumentModel } from "../model";
+import { initialDocument } from "./protocol-fixtures";
+import { createTestNotebook } from "./test-notebook";
+import { createUser } from "./user";
+
+function createSession(initial: Document, controller: AbortController) {
+  const store = createStore();
+  const { notebook, control } = createTestNotebook(store, initial);
+  const model = createDocumentModel({
+    store,
+    documentAtom: notebook.documentAtom,
+    backendAtom: notebook.backendAtom,
+  });
+  model.start({ signal: controller.signal });
+  return {
+    user: createUser(store, model),
+    host: control,
+    notebook,
+    controller,
+  };
+}
+
+export const test = baseTest.extend<{
+  initial: Document;
+  session: ReturnType<typeof createSession>;
+}>({
+  initial: initialDocument(),
+  session: async ({ initial }, provide) => {
+    vi.useFakeTimers({ now: 0 });
+    const controller = new AbortController();
+    try {
+      await provide(createSession(initial, controller));
+    } finally {
+      controller.abort();
+      vi.useRealTimers();
+    }
+  },
+});

package/src/core/notebook/document/__tests__/scenarios.test.ts

@@ -0,0 +1,205 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { vi } from "vitest";
+import { test } from "./scenario-fixture";
+
+test("waits for a pause in typing before sending the latest edit", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "first");
+  await vi.advanceTimersByTimeAsync(300);
+  user.changeCode("c3", "latest");
+  await vi.advanceTimersByTimeAsync(399);
+  expect(notebook.edit).not.toHaveBeenCalled();
+  await vi.advanceTimersByTimeAsync(1);
+  expect(notebook.edit).toHaveBeenCalledOnce();
+  const request = await host.nextRequest();
+  request.process();
+  request.deliverResponse();
+  host.deliverDocument();
+  await expect(user.sync()).resolves.toEqual(host.document().revision);
+  expect(host.document().cells.get("c3")?.code).toBe("latest");
+});
+
+test.for(["response", "document"] as const)(
+  "preserves newer typing when the %s arrives first",
+  async (first, { expect, session: { user, host, notebook } }) => {
+    user.changeCells({ c3: { code: "sent", config: { hide_code: true } } });
+    const flushed = user.sync();
+    const finished = vi.fn();
+    void flushed.then(finished, () => undefined);
+    const request = await host.nextRequest();
+    request.process();
+    user.changeCode("c3", "still typing");
+
+    const deliver = {
+      response: () => request.deliverResponse(),
+      document: () => host.deliverDocument(),
+    };
+    deliver[first]();
+    await vi.advanceTimersByTimeAsync(0);
+    expect(finished).not.toHaveBeenCalled();
+    expect(notebook.edit).toHaveBeenCalledOnce();
+    expect(user.cell("c3")?.code).toBe("still typing");
+
+    deliver[first === "response" ? "document" : "response"]();
+    await vi.advanceTimersByTimeAsync(0);
+    expect(user.cell("c3")?.code).toBe("still typing");
+    const next = await host.nextRequest();
+    next.process();
+    next.deliverResponse();
+    host.deliverDocument();
+    await expect(flushed).resolves.toEqual(host.document().revision);
+    const expected = {
+      id: "c3",
+      code: "still typing",
+      name: "_",
+      config: {
+        column: null,
+        disabled: false,
+        hide_code: true,
+        expand_output: false,
+      },
+    };
+    expect(user.cell("c3")).toEqual(expected);
+    expect(host.document().cells.get("c3")).toMatchObject(expected);
+    expect(user.status()).toEqual({ status: "idle" });
+
+    host.changeCell("c3", {
+      code: "another writer",
+      config: { hide_code: false },
+    });
+    host.deliverDocument();
+    expect(user.cell("c3")?.code).toBe("another writer");
+    expect(user.cell("c3")?.config.hide_code).toBe(false);
+  },
+);
+
+test("uses a newer confirmed document for edits that have not been sent", async ({
+  expect,
+  session: { user, host },
+}) => {
+  user.changeCode("c3", "local");
+  host.changeCell("c3", { name: "remote name" });
+  host.deliverDocument();
+  const flushed = user.sync();
+  const request = await host.nextRequest();
+  request.process();
+  request.deliverResponse();
+  host.deliverDocument();
+  await expect(flushed).resolves.toEqual(host.document().revision);
+  expect(host.document().cells.get("c3")).toMatchObject({
+    code: "local",
+    name: "remote name",
+  });
+});
+
+test.for(["before reset", "during reset", "after snapshot"] as const)(
+  "recovers across reset when the response arrives %s",
+  async (arrival, { expect, session: { user, host } }) => {
+    user.changeCode("c3", "sent");
+    const flushed = user.sync();
+    const request = await host.nextRequest();
+    request.process();
+    user.changeCode("c3", "newer typing");
+    if (arrival === "before reset") {
+      request.deliverResponse();
+      await vi.advanceTimersByTimeAsync(0);
+    }
+    host.reset();
+    if (arrival === "during reset") {
+      request.deliverResponse();
+      await vi.advanceTimersByTimeAsync(0);
+    }
+    host.changeCell("c1", { code: "unrelated edit" });
+    host.deliverSnapshot();
+    if (arrival === "after snapshot") {
+      await vi.advanceTimersByTimeAsync(0);
+      request.deliverResponse();
+    }
+    expect(user.cell("c3")?.code).toBe("newer typing");
+    const next = await host.nextRequest();
+    next.process();
+    next.deliverResponse();
+    host.deliverDocument();
+    await expect(flushed).resolves.toEqual(host.document().revision);
+    expect(host.document().cells.get("c3")?.code).toBe("newer typing");
+    expect(user.cell("c1")?.code).toBe("unrelated edit");
+    expect(user.status()).toEqual({ status: "idle" });
+
+    host.changeCell("c3", { code: "another writer" });
+    host.deliverDocument();
+    expect(user.cell("c3")?.code).toBe("another writer");
+  },
+);
+
+test.for(["response", "document"] as const)(
+  "does not resurrect a cancelled creation when the %s arrives first",
+  async (first, { expect, session: { user, host } }) => {
+    user.createCell({ id: "new", code: "new cell" });
+    const flushed = user.sync();
+    const creation = await host.nextRequest();
+    user.deleteCell("new");
+    host.changeCell("c1", { code: "unrelated edit" });
+    host.deliverDocument();
+    creation.process();
+    const deliver = {
+      response: () => creation.deliverResponse(),
+      document: () => host.deliverDocument(),
+    };
+    deliver[first]();
+    await vi.advanceTimersByTimeAsync(0);
+    expect(user.cell("new")).toBeUndefined();
+    expect(user.order()).toEqual(["c1", "c3", "k7f2"]);
+    deliver[first === "response" ? "document" : "response"]();
+    const deletion = await host.nextRequest();
+    deletion.process();
+    deletion.deliverResponse();
+    host.deliverDocument();
+    await expect(flushed).resolves.toEqual(host.document().revision);
+    expect(host.document().cells.has("new")).toBe(false);
+    expect(user.cell("new")).toBeUndefined();
+  },
+);
+
+test("releases an edit undone before sending and follows subsequent host changes", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  user.changeCode("c3", "temporary");
+  user.changeCode("c3", "x = 21");
+  expect(user.status()).toEqual({ status: "idle" });
+  await expect(user.sync()).resolves.toEqual(host.document().revision);
+
+  host.changeCell("c3", { code: "another writer" });
+  host.deliverDocument();
+  expect(user.cell("c3")?.code).toBe("another writer");
+  await vi.advanceTimersByTimeAsync(1000);
+  expect(notebook.edit).not.toHaveBeenCalled();
+});
+
+test("flushes edits made while an earlier flush is finishing", async ({
+  expect,
+  session: { user, host, notebook },
+}) => {
+  const original = host.document().revision;
+  const first = user.sync();
+  // Let the first flush drain without awaiting its returned promise.
+  await Promise.resolve();
+  user.changeCode("c3", "new code");
+  const second = user.sync();
+  const settled = vi.fn();
+  void second.then(settled, settled);
+
+  await vi.advanceTimersByTimeAsync(0);
+  expect(notebook.edit).toHaveBeenCalledOnce();
+  expect(settled).not.toHaveBeenCalled();
+  const request = await host.nextRequest();
+  request.process();
+  expect(host.document().cells.get("c3")?.code).toBe("new code");
+  request.deliverResponse();
+  host.deliverDocument();
+  await expect(second).resolves.toEqual(host.document().revision);
+  await expect(first).resolves.toEqual(original);
+});