@marimo-team/islands 0.25.3-dev19 → 0.25.3-dev191

package/src/core/notebook/document/sync.ts

@@ -0,0 +1,418 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { idempotencyKey, type Result } from "@marimo-team/client";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import type { Document, NotebookBackend } from "../replica";
+import type { DocumentState, DocumentSubmission } from "./state";
+
+export type DocumentSyncStatus =
+  | {
+      readonly status:
+        | "idle"
+        | "pending"
+        | "sending"
+        | "confirming"
+        | "unavailable";
+    }
+  | {
+      readonly status: "error";
+      readonly error: Error;
+      readonly problem?: Protocol.Problem;
+      readonly retryable: boolean;
+    };
+
+interface Submission {
+  readonly submission: DocumentSubmission;
+  /** The backend that confirmed the submission's base. */
+  readonly backend: NotebookBackend;
+  readonly idempotencyKey: string;
+  readonly startedAt: number;
+  readonly controller: AbortController;
+  response?: Protocol.Document;
+  problem?: Protocol.Problem;
+}
+
+// The marimo host retains receipts for five minutes. Stop well before that,
+// including after sleep, rather than replaying an uncertain write as a new one.
+const RETRY_WINDOW_MS = 60_000;
+const REQUEST_TIMEOUT_MS = 10_000;
+const MAX_ATTEMPTS = 3;
+const DEBOUNCE_MS = 400;
+
+function abortable<T>(promise: Promise<T>, signal: AbortSignal): Promise<T> {
+  return new Promise((resolve, reject) => {
+    const abort = () => reject(signal.reason);
+    signal.addEventListener("abort", abort, { once: true });
+    promise
+      .then(resolve, reject)
+      .finally(() => signal.removeEventListener("abort", abort));
+    if (signal.aborted) {
+      abort();
+    }
+  });
+}
+
+function delay(ms: number, signal: AbortSignal): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const abort = () => {
+      clearTimeout(timer);
+      reject(signal.reason);
+    };
+    const timer = setTimeout(() => {
+      signal.removeEventListener("abort", abort);
+      resolve();
+    }, ms);
+    signal.addEventListener("abort", abort, { once: true });
+    if (signal.aborted) {
+      abort();
+    }
+  });
+}
+
+/** Own submission scheduling; follow the existing source without opening a stream. */
+export function createDocumentSync({
+  store,
+  documentAtom: source,
+  backendAtom,
+  state,
+  signal,
+}: {
+  store: ReturnType<typeof createStore>;
+  documentAtom: Atom<Document | null>;
+  backendAtom: Atom<NotebookBackend | null>;
+  state: DocumentState;
+  signal: AbortSignal;
+}) {
+  signal.throwIfAborted();
+  const statusAtom = atom<DocumentSyncStatus>({ status: "idle" });
+  const observed = atom((get) => ({
+    document: get(source),
+    backend: get(backendAtom),
+    pending: get(state.pendingAtom),
+    changes: get(state.changesAtom),
+  }));
+  let active: Submission | null = null;
+  let running: Promise<Protocol.Document> | null = null;
+  let failure: Extract<DocumentSyncStatus, { status: "error" }> | null = null;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  let refreshing = false;
+  // Holds keep edits local, as a legacy save needs while it writes the
+  // notebook the flush before it produced.
+  let holds = 0;
+  let releaseWaiters: Array<{
+    resolve: () => void;
+    reject: (reason: unknown) => void;
+  }> = [];
+
+  function clearTimer() {
+    clearTimeout(timer);
+    timer = undefined;
+  }
+
+  function unavailable() {
+    return new Error(
+      "The document or source required by an edit is unavailable",
+    );
+  }
+
+  function refresh() {
+    if (signal.aborted || refreshing) {
+      return;
+    }
+    refreshing = true;
+    try {
+      const current = store.get(source);
+      if (
+        active &&
+        (store.get(backendAtom) !== active.backend ||
+          (current &&
+            current.revision.notebook_id !== active.submission.notebookId))
+      ) {
+        active.controller.abort(
+          new Error("The notebook changed during synchronization"),
+        );
+      }
+      if (running || failure) {
+        return;
+      }
+      clearTimer();
+      const changes = store.get(state.changesAtom);
+      if (changes === null || store.get(backendAtom) === null) {
+        store.set(statusAtom, { status: "unavailable" });
+      } else if (changes.length === 0) {
+        // Release unsent reversals instead of keeping an invisible override.
+        store.set(state.submitAtom);
+        store.set(statusAtom, { status: "idle" });
+      } else {
+        store.set(statusAtom, { status: "pending" });
+        if (holds === 0) {
+          timer = setTimeout(() => {
+            void flush().catch(() => undefined);
+          }, DEBOUNCE_MS);
+        }
+      }
+    } finally {
+      refreshing = false;
+    }
+  }
+
+  async function send(request: Submission) {
+    for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
+      signal.throwIfAborted();
+      request.controller.signal.throwIfAborted();
+      if (Date.now() - request.startedAt >= RETRY_WINDOW_MS) {
+        throw new Error(
+          "The edit's retry window expired; its outcome is unknown",
+        );
+      }
+      store.set(statusAtom, { status: "sending" });
+      const controller = new AbortController();
+      const abort = () => controller.abort(request.controller.signal.reason);
+      request.controller.signal.addEventListener("abort", abort, {
+        once: true,
+      });
+      const timeout = setTimeout(
+        () => controller.abort(new Error("The edit request timed out")),
+        REQUEST_TIMEOUT_MS,
+      );
+      let result: Result<Protocol.EditResponseOutput> | undefined;
+      try {
+        signal.throwIfAborted();
+        request.controller.signal.throwIfAborted();
+        result = await abortable(
+          request.backend
+            .notebook(request.submission.notebookId)
+            .edit(request.submission.changes, {
+              idempotencyKey: request.idempotencyKey,
+              signal: controller.signal,
+            }),
+          controller.signal,
+        );
+      } catch (error) {
+        // The host may or may not have applied it. Repeat it with the same key.
+        if (request.controller.signal.aborted || attempt === MAX_ATTEMPTS - 1) {
+          throw error;
+        }
+      } finally {
+        clearTimeout(timeout);
+        request.controller.signal.removeEventListener("abort", abort);
+      }
+      signal.throwIfAborted();
+      request.controller.signal.throwIfAborted();
+      if (result === undefined) {
+        await delay(500 * 2 ** attempt, request.controller.signal);
+        continue;
+      }
+      // The host answered, so repeating the request returns the same answer.
+      if (!result.ok) {
+        request.problem = result.problem;
+        throw new Error(result.problem.detail ?? result.problem.title, {
+          cause: result.problem,
+        });
+      }
+      if (result.data.document.notebook_id !== request.submission.notebookId) {
+        throw new Error("Edit response belongs to another notebook");
+      }
+      request.response = result.data.document;
+      store.set(state.settleAtom, request.submission, request.response);
+      return;
+    }
+  }
+
+  function confirm(request: Submission): Promise<void> {
+    store.set(statusAtom, { status: "confirming" });
+    return new Promise((resolve, reject) => {
+      const check = () => {
+        if (signal.aborted || request.controller.signal.aborted) {
+          cleanup();
+          reject(
+            signal.aborted ? signal.reason : request.controller.signal.reason,
+          );
+        } else if (store.get(state.pendingAtom) !== request.submission) {
+          cleanup();
+          resolve();
+        }
+      };
+      const unsubscribe = store.sub(observed, check);
+      const cleanup = () => {
+        unsubscribe();
+        request.controller.signal.removeEventListener("abort", check);
+      };
+      request.controller.signal.addEventListener("abort", check, {
+        once: true,
+      });
+      check();
+    });
+  }
+
+  async function drain(
+    notebookId: Protocol.NotebookId,
+  ): Promise<Protocol.Document> {
+    try {
+      while (true) {
+        signal.throwIfAborted();
+        const current = store.get(source);
+        if (current && current.revision.notebook_id !== notebookId) {
+          throw new Error("The notebook changed during synchronization");
+        }
+        if (active === null) {
+          const backend = store.get(backendAtom);
+          if (
+            current === null ||
+            backend === null ||
+            store.get(state.changesAtom) === null
+          ) {
+            throw unavailable();
+          }
+          const submission = store.set(state.submitAtom);
+          if (submission === null) {
+            return current.revision;
+          }
+          active = {
+            submission,
+            backend,
+            idempotencyKey: idempotencyKey(),
+            startedAt: Date.now(),
+            controller: new AbortController(),
+          };
+        }
+        if (!active.response) {
+          await send(active);
+        }
+        await confirm(active);
+        active = null;
+      }
+    } catch (cause: unknown) {
+      signal.throwIfAborted();
+      const error =
+        cause instanceof Error
+          ? cause
+          : new Error("Document synchronization failed", { cause });
+      failure = {
+        status: "error",
+        error,
+        problem: active?.problem,
+        retryable:
+          !signal.aborted &&
+          !active?.controller.signal.aborted &&
+          (active === null ||
+            active.problem !== undefined ||
+            Date.now() - active.startedAt < RETRY_WINDOW_MS),
+      };
+      if (!signal.aborted) {
+        store.set(statusAtom, failure);
+      }
+      throw error;
+    } finally {
+      // Stop accepting joiners in the same turn that the drain finishes.
+      running = null;
+      refresh();
+    }
+  }
+
+  /** Drain edits and return the revision to use as the subsequent command's precondition. */
+  function flush(): Promise<Protocol.Document> {
+    clearTimer();
+    if (signal.aborted) {
+      return Promise.reject(signal.reason);
+    }
+    if (failure) {
+      return Promise.reject(failure.error);
+    }
+    if (running) {
+      return running;
+    }
+    if (holds > 0) {
+      return new Promise<void>((resolve, reject) => {
+        releaseWaiters.push({ resolve, reject });
+      }).then(flush);
+    }
+    const notebookId =
+      active?.submission.notebookId ?? store.get(source)?.revision.notebook_id;
+    if (notebookId === undefined) {
+      return Promise.reject(unavailable());
+    }
+    running = Promise.resolve().then(() => drain(notebookId));
+    return running;
+  }
+
+  /** Explicitly retry a rejection against fresh state, or repeat an uncertain request exactly. */
+  function retry(): Promise<Protocol.Document> {
+    if (signal.aborted) {
+      return Promise.reject(signal.reason);
+    }
+    if (!failure) {
+      return flush();
+    }
+    if (!failure.retryable) {
+      return Promise.reject(failure.error);
+    }
+    if (active?.problem) {
+      const current = store.get(source);
+      const rejectedAt = active.problem.edit?.document.version;
+      if (
+        !current ||
+        current.revision.notebook_id !== active.submission.notebookId ||
+        (rejectedAt !== undefined && current.revision.version < rejectedAt)
+      ) {
+        return Promise.reject(
+          new Error("Wait for the host's current document before retrying"),
+        );
+      }
+      store.set(state.settleAtom, active.submission, null);
+      active = null;
+    }
+    failure = null;
+    return flush();
+  }
+
+  /**
+   * Keep edits local until the returned release is called. Flushes
+   * requested meanwhile wait for it. Holds nest, and a release is
+   * idempotent.
+   *
+   * TODO(MO-8028): Only the legacy save needs this, since the host takes
+   * the notebook it writes as the document. Saving through the host
+   * removes both.
+   */
+  function hold(): () => void {
+    holds += 1;
+    let held = true;
+    return () => {
+      if (!held) {
+        return;
+      }
+      held = false;
+      holds -= 1;
+      if (holds > 0) {
+        return;
+      }
+      const waiters = releaseWaiters;
+      releaseWaiters = [];
+      for (const waiter of waiters) {
+        waiter.resolve();
+      }
+      refresh();
+    };
+  }
+
+  const unsubscribe = store.sub(observed, refresh);
+  function stop() {
+    clearTimer();
+    unsubscribe();
+    signal.removeEventListener("abort", stop);
+    active?.controller.abort(signal.reason);
+    const waiters = releaseWaiters;
+    releaseWaiters = [];
+    for (const waiter of waiters) {
+      waiter.reject(signal.reason);
+    }
+  }
+  signal.addEventListener("abort", stop, { once: true });
+  if (signal.aborted) {
+    stop();
+  }
+  refresh();
+  return { statusAtom, flush, retry, hold };
+}

package/src/core/notebook/index.ts

@@ -0,0 +1,18 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** React hooks for the current notebook. */
+
+export {
+  useNotebookAttachments,
+  attachmentLabel,
+  type Attachment,
+} from "./attachments";
+export {
+  useNotebookCellCode,
+  useNotebookCellCodeExtension,
+  useNotebookCellConfig,
+  useNotebookCellName,
+  useNotebookCellOrder,
+  useNotebookDocumentEdit,
+  useNotebookDocumentSync,
+} from "./document";

package/src/core/notebook/replica/AGENTS.md

@@ -0,0 +1,11 @@
+# Notebook replica
+
+Production code may import replica modules, `@marimo-team/client`,
+`jotai/vanilla`, `zod`, and `@/utils/Logger`. Enforced by
+`marimo/replica-imports`. Tests are exempt.
+
+Code outside this directory imports it as `@/core/notebook/replica`.
+Enforced by `marimo/module-entry`.
+
+The protocol schemas and types live in `@marimo-team/client/protocol`.
+Define interfaces here and inject application behavior through adapters.

package/src/core/notebook/replica/__tests__/attachments.test.ts

@@ -0,0 +1,56 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { attachments } from "../resources/attachments";
+import { createNotebookReducer } from "@marimo-team/client";
+import { fixtureEvents } from "./fixtures";
+
+const reducer = createNotebookReducer({ attachments });
+const snapshot = fixtureEvents("notebook-events.sse");
+const records = snapshot.flatMap((event) =>
+  event.event === "attachment" ? [event.data.attachment] : [],
+);
+const confirmed = () => snapshot.reduce(reducer.reduce, reducer.initial);
+
+describe("attachment resource", () => {
+  it("collects snapshot records, replaces by id, and removes attachments", () => {
+    const state = confirmed();
+    expect(state.resources.attachments).toEqual(records);
+    const [browser, agent] = records;
+    const renamed = { ...agent, name: "Renamed agent" };
+    const updated = reducer.reduce(state, {
+      event: "attachment",
+      id: "98",
+      data: { attachment: renamed },
+    });
+    expect(updated.resources.attachments).toEqual([browser, renamed]);
+    expect(state.resources.attachments).toEqual(records);
+
+    const removed = reducer.reduce(updated, {
+      event: "attachment.removed",
+      id: "99",
+      data: { id: browser.id, time: "2026-10-01T09:30:00Z" },
+    });
+    expect(removed.resources.attachments).toEqual([renamed]);
+    const repeated = reducer.reduce(removed, {
+      event: "attachment.removed",
+      id: "100",
+      data: { id: browser.id, time: "2026-10-01T09:30:00Z" },
+    });
+    expect(repeated.resources.attachments).toBe(removed.resources.attachments);
+  });
+
+  it("rebuilds records after reset", () => {
+    const state = confirmed();
+    const reset = reducer.reduce(state, { event: "reset", data: {} });
+    expect(reset.resources.attachments).toEqual([]);
+    const replacement = snapshot
+      .filter(
+        (event) =>
+          event.event !== "attachment" ||
+          event.data.attachment.id === records[0].id,
+      )
+      .reduce(reducer.reduce, reset);
+    expect(replacement.resources.attachments).toEqual([records[0]]);
+  });
+});

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

@@ -0,0 +1,126 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { document } from "../resources/document";
+import { createNotebookReducer } from "@marimo-team/client";
+import { fixtureEvents } from "./fixtures";
+
+const reducer = createNotebookReducer({ document });
+const snapshot = fixtureEvents("notebook-events.sse");
+const edits = fixtureEvents("edits.sse");
+
+function replay(
+  events: readonly Protocol.EventOutput[],
+  from = reducer.initial,
+) {
+  return events.reduce((state, event) => reducer.reduce(state, event), from);
+}
+
+const confirmed = () => replay(snapshot);
+const current = (state: ReturnType<typeof replay>) =>
+  state.resources.document.current;
+
+describe("document resource", () => {
+  it("publishes the snapshot at ready, then each completed edit and save", () => {
+    const ready = snapshot.findIndex((event) => event.event === "ready");
+    const opening = replay(snapshot.slice(0, ready));
+    expect(current(opening)).toBeNull();
+    const first = current(replay(snapshot.slice(ready), opening));
+    expect(first?.revision.order).toEqual(["c1", "c3", "k7f2"]);
+    expect([...first!.cells.keys()]).toEqual(first?.revision.order);
+
+    const [edit, save] = edits.filter((event) => event.event === "document");
+    const pending = replay(
+      edits.filter((event) => event.event !== "document"),
+      confirmed(),
+    );
+    expect(current(pending)).toEqual(first);
+    const changed = reducer.reduce(pending, edit);
+    expect(current(changed)?.revision).toEqual(edit.data.document);
+    expect([...current(changed)!.cells.keys()].toSorted()).toEqual(
+      edit.data.document.order.toSorted(),
+    );
+    const saved = current(reducer.reduce(changed, save));
+    expect(saved?.revision).toEqual(save.data.document);
+    expect(saved?.cells.get("c3")?.code).toBe("x = 22");
+    expect(saved?.cells.has("k7f2")).toBe(false);
+    expect(saved?.revision.saved_etag).toBe(saved?.revision.etag);
+    expect(first?.cells.get("c3")?.code).toBe("x = 21");
+    expect(first?.cells.has("k7f2")).toBe(true);
+  });
+
+  it("applies order-only changes and distinguishes an empty notebook from no document", () => {
+    const state = confirmed();
+    const revision = current(state)!.revision;
+    const reordered = {
+      ...revision,
+      order: revision.order.toReversed(),
+      version: revision.version + 1,
+    };
+    const changed = reducer.reduce(state, {
+      event: "document",
+      data: { document: reordered },
+      id: "200",
+    });
+    expect(current(changed)).toEqual({
+      revision: reordered,
+      cells: current(state)?.cells,
+    });
+
+    const empty = { ...revision, order: [] };
+    const opening = reducer.reduce(reducer.initial, {
+      event: "document",
+      data: { document: empty },
+    });
+    expect(current(opening)).toBeNull();
+    expect(
+      current(reducer.reduce(opening, { event: "ready", data: {}, id: "0" })),
+    ).toEqual({ revision: empty, cells: new Map() });
+    expect(
+      current(
+        reducer.reduce(reducer.initial, { event: "ready", data: {}, id: "0" }),
+      ),
+    ).toBeNull();
+  });
+
+  it("keeps the document while a live update is pending and after a replay completes it", () => {
+    const state = confirmed();
+    const [cell, , removed, published] = edits;
+    if (published.event !== "document") {
+      throw new Error("Fixture does not publish a document");
+    }
+    const partial = replay([cell, removed], state);
+    expect(current(partial)).toBe(current(state));
+    const completed = replay(
+      [published, { event: "ready", data: {}, id: "113" }],
+      partial,
+    );
+    expect(current(completed)?.revision).toEqual(published.data.document);
+  });
+
+  it("removes a notebook without waiting for individual cell removals", () => {
+    const removed = reducer.reduce(confirmed(), {
+      event: "notebook.removed",
+      data: { id: "n1", time: "2026-10-01T09:30:00Z" },
+      id: "200",
+    });
+    expect(current(removed)).toBeNull();
+  });
+
+  it("rejects records that cannot form a complete revision", () => {
+    const state = confirmed();
+    const revision = current(state)!.revision;
+    expect(() =>
+      reducer.reduce(state, {
+        event: "document",
+        data: { document: { ...revision, order: ["missing"] } },
+        id: "200",
+      }),
+    ).toThrow("Document records do not match");
+    const partial = reducer.reduce(state, edits[0]);
+    expect(() =>
+      reducer.reduce(partial, { event: "ready", data: {}, id: "200" }),
+    ).toThrow("incomplete document update");
+  });
+});