@marimo-team/islands 0.25.3-dev29 → 0.25.3-dev291

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

@@ -0,0 +1,369 @@
+/* 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;
+
+  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" });
+        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;
+    }
+    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();
+  }
+
+  const unsubscribe = store.sub(observed, refresh);
+  function stop() {
+    clearTimer();
+    unsubscribe();
+    signal.removeEventListener("abort", stop);
+    active?.controller.abort(signal.reason);
+  }
+  signal.addEventListener("abort", stop, { once: true });
+  if (signal.aborted) {
+    stop();
+  }
+  refresh();
+  return { statusAtom, flush, retry };
+}

package/src/core/notebook/index.ts

@@ -0,0 +1,19 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What the editor reads about the current notebook. */
+
+export {
+  useNotebookAttachments,
+  attachmentLabel,
+  type Attachment,
+} from "./attachments";
+export {
+  useNotebookCellCode,
+  useNotebookCellCodeExtension,
+  useNotebookCellConfig,
+  useNotebookCellName,
+  useNotebookCellOrder,
+  useNotebookDocumentEdit,
+  useNotebookDocumentSync,
+} from "./document";
+export { canEditAtom, notebookOperationsAtom } from "./operations";

package/src/core/notebook/operations.ts

@@ -0,0 +1,30 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What the host lets this client do to the notebook. */
+
+import { atom } from "jotai";
+import { kioskModeAtom } from "../mode";
+import { activeNotebookReplicaAtom } from "../session/active-notebook";
+
+const NO_OPERATIONS: readonly string[] = [];
+
+/** Operations the host advertises for this notebook, none until its stream is ready. */
+export const notebookOperationsAtom = atom((get) => {
+  const replica = get(activeNotebookReplicaAtom);
+  const notebook = replica === null ? null : get(replica.notebookAtom);
+  return notebook?.operations ?? NO_OPERATIONS;
+});
+
+/**
+ * Whether this client may edit the document.
+ *
+ * TODO(MO-8352): The host advertises the same operations to every
+ * connection of a notebook, so a kiosk connection, an interactor that may
+ * not edit, narrows it here. Once the host narrows `operations` per
+ * connection by consumer capability, the kiosk term goes.
+ */
+export const canEditAtom = atom(
+  (get) =>
+    !get(kioskModeAtom) &&
+    get(notebookOperationsAtom).includes("notebook.edit"),
+);

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/{host → notebook/replica}/__tests__/attachments.test.ts

@@ -2,15 +2,15 @@
 
 import { describe, expect, it } from "vitest";
 import { attachments } from "../resources/attachments";
-import { createStream } from "../stream";
+import { createNotebookReducer } from "@marimo-team/client";
 import { fixtureEvents } from "./fixtures";
 
-const stream = createStream({ attachments });
+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(stream.reduce, stream.initial);
+const confirmed = () => snapshot.reduce(reducer.reduce, reducer.initial);
 
 describe("attachment resource", () => {
   it("collects snapshot records, replaces by id, and removes attachments", () => {
@@ -18,7 +18,7 @@ describe("attachment resource", () => {
     expect(state.resources.attachments).toEqual(records);
     const [browser, agent] = records;
     const renamed = { ...agent, name: "Renamed agent" };
-    const updated = stream.reduce(state, {
+    const updated = reducer.reduce(state, {
       event: "attachment",
       id: "98",
       data: { attachment: renamed },
@@ -26,13 +26,13 @@ describe("attachment resource", () => {
     expect(updated.resources.attachments).toEqual([browser, renamed]);
     expect(state.resources.attachments).toEqual(records);
 
-    const removed = stream.reduce(updated, {
+    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 = stream.reduce(removed, {
+    const repeated = reducer.reduce(removed, {
       event: "attachment.removed",
       id: "100",
       data: { id: browser.id, time: "2026-10-01T09:30:00Z" },
@@ -40,10 +40,9 @@ describe("attachment resource", () => {
     expect(repeated.resources.attachments).toBe(removed.resources.attachments);
   });
 
-  it("retains records for replay and rebuilds them after reset", () => {
+  it("rebuilds records after reset", () => {
     const state = confirmed();
-    expect(stream.resume(state)).toBe(state);
-    const reset = stream.reduce(state, { event: "reset", data: {} });
+    const reset = reducer.reduce(state, { event: "reset", data: {} });
     expect(reset.resources.attachments).toEqual([]);
     const replacement = snapshot
       .filter(
@@ -51,7 +50,7 @@ describe("attachment resource", () => {
           event.event !== "attachment" ||
           event.data.attachment.id === records[0].id,
       )
-      .reduce(stream.reduce, reset);
+      .reduce(reducer.reduce, reset);
     expect(replacement.resources.attachments).toEqual([records[0]]);
   });
 });

package/src/core/{host → notebook/replica}/__tests__/document.test.ts

@@ -1,17 +1,20 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { describe, expect, it } from "vitest";
-import type { NotebookEvent } from "../events";
+import type * as Protocol from "@marimo-team/client/protocol";
 import { document } from "../resources/document";
-import { createStream } from "../stream";
+import { createNotebookReducer } from "@marimo-team/client";
 import { fixtureEvents } from "./fixtures";
 
-const stream = createStream({ document });
+const reducer = createNotebookReducer({ document });
 const snapshot = fixtureEvents("notebook-events.sse");
 const edits = fixtureEvents("edits.sse");
 
-function replay(events: readonly NotebookEvent[], from = stream.initial) {
-  return events.reduce((state, event) => stream.reduce(state, event), from);
+function replay(
+  events: readonly Protocol.EventOutput[],
+  from = reducer.initial,
+) {
+  return events.reduce((state, event) => reducer.reduce(state, event), from);
 }
 
 const confirmed = () => replay(snapshot);
@@ -33,12 +36,12 @@ describe("document resource", () => {
       confirmed(),
     );
     expect(current(pending)).toEqual(first);
-    const changed = stream.reduce(pending, edit);
+    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(stream.reduce(changed, save));
+    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);
@@ -55,7 +58,7 @@ describe("document resource", () => {
       order: revision.order.toReversed(),
       version: revision.version + 1,
     };
-    const changed = stream.reduce(state, {
+    const changed = reducer.reduce(state, {
       event: "document",
       data: { document: reordered },
       id: "200",
@@ -66,17 +69,17 @@ describe("document resource", () => {
     });
 
     const empty = { ...revision, order: [] };
-    const opening = stream.reduce(stream.initial, {
+    const opening = reducer.reduce(reducer.initial, {
       event: "document",
       data: { document: empty },
     });
     expect(current(opening)).toBeNull();
     expect(
-      current(stream.reduce(opening, { event: "ready", data: {}, id: "0" })),
+      current(reducer.reduce(opening, { event: "ready", data: {}, id: "0" })),
     ).toEqual({ revision: empty, cells: new Map() });
     expect(
       current(
-        stream.reduce(stream.initial, { event: "ready", data: {}, id: "0" }),
+        reducer.reduce(reducer.initial, { event: "ready", data: {}, id: "0" }),
       ),
     ).toBeNull();
   });
@@ -89,7 +92,6 @@ describe("document resource", () => {
     }
     const partial = replay([cell, removed], state);
     expect(current(partial)).toBe(current(state));
-    expect(stream.resume(partial)).toBe(partial);
     const completed = replay(
       [published, { event: "ready", data: {}, id: "113" }],
       partial,
@@ -98,7 +100,7 @@ describe("document resource", () => {
   });
 
   it("removes a notebook without waiting for individual cell removals", () => {
-    const removed = stream.reduce(confirmed(), {
+    const removed = reducer.reduce(confirmed(), {
       event: "notebook.removed",
       data: { id: "n1", time: "2026-10-01T09:30:00Z" },
       id: "200",
@@ -110,15 +112,15 @@ describe("document resource", () => {
     const state = confirmed();
     const revision = current(state)!.revision;
     expect(() =>
-      stream.reduce(state, {
+      reducer.reduce(state, {
         event: "document",
         data: { document: { ...revision, order: ["missing"] } },
         id: "200",
       }),
     ).toThrow("Document records do not match");
-    const partial = stream.reduce(state, edits[0]);
+    const partial = reducer.reduce(state, edits[0]);
     expect(() =>
-      stream.reduce(partial, { event: "ready", data: {}, id: "200" }),
+      reducer.reduce(partial, { event: "ready", data: {}, id: "200" }),
     ).toThrow("incomplete document update");
   });
 });