@marimo-team/frontend 0.25.3-dev10 → 0.25.3-dev100

package/src/core/document/sync.ts

@@ -0,0 +1,359 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { idempotencyKey, type Result } from "@marimo-team/client";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, type createStore } from "jotai/vanilla";
+import type { NotebookClient } from "../host/notebook-client";
+import type { DocumentState, DocumentSubmission } from "./state";
+
+/** The notebook state and operation needed by document editing. */
+export type DocumentHost = Pick<NotebookClient, "documentAtom" | "edit">;
+
+export type DocumentSyncState =
+  | {
+      readonly status:
+        | "idle"
+        | "pending"
+        | "sending"
+        | "confirming"
+        | "unavailable";
+    }
+  | {
+      readonly status: "error";
+      readonly error: Error;
+      readonly problem?: Protocol.Problem;
+      readonly retryable: boolean;
+    };
+
+interface Submission {
+  readonly edit: DocumentSubmission;
+  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,
+  notebook,
+  state,
+  signal,
+}: {
+  store: ReturnType<typeof createStore>;
+  notebook: DocumentHost;
+  state: DocumentState;
+  signal: AbortSignal;
+}) {
+  signal.throwIfAborted();
+  const source = notebook.documentAtom;
+  const stateAtom = atom<DocumentSyncState>({ status: "idle" });
+  const observed = atom((get) => ({
+    document: get(source),
+    pending: get(state.pendingAtom),
+    changes: get(state.changesAtom),
+  }));
+  let active: Submission | null = null;
+  let running: Promise<Protocol.Document> | null = null;
+  let failure: Extract<DocumentSyncState, { 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 &&
+        current &&
+        current.revision.notebook_id !== active.edit.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.set(stateAtom, { status: "unavailable" });
+      } else if (changes.length === 0) {
+        // Release unsent reversals instead of keeping an invisible override.
+        store.set(state.submitAtom);
+        store.set(stateAtom, { status: "idle" });
+      } else {
+        store.set(stateAtom, { 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(stateAtom, { 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(
+          notebook.edit(request.edit.notebookId, request.edit.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.edit.notebookId) {
+        throw new Error("Edit response belongs to another notebook");
+      }
+      request.response = result.data.document;
+      store.set(state.settleAtom, request.edit, request.response);
+      return;
+    }
+  }
+
+  function confirm(request: Submission): Promise<void> {
+    store.set(stateAtom, { 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.edit) {
+          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) {
+          if (current === null || store.get(state.changesAtom) === null) {
+            throw unavailable();
+          }
+          const edit = store.set(state.submitAtom);
+          if (edit === null) {
+            return current.revision;
+          }
+          active = {
+            edit,
+            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(stateAtom, 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?.edit.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.edit.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.edit, 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 { stateAtom, flush, retry };
+}

package/src/core/edit-app.tsx

@@ -41,7 +41,7 @@ import { viewStateAtom } from "./mode";
 import { useRequestClient } from "./network/requests";
 import { useFilename } from "./saving/filename";
 import { lastSavedNotebookAtom } from "./saving/state";
-import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
+import { useNotebookSession } from "./websocket/useNotebookSession";
 
 interface AppProps {
   /**
@@ -83,7 +83,7 @@ export const EditApp: React.FC<AppProps> = ({
     };
   }, []);
 
-  const { connection, reconnect } = useMarimoKernelConnection({
+  const { connection, reconnect } = useNotebookSession({
     autoInstantiate: userConfig.runtime.auto_instantiate,
     setCells: (cells, layout) => {
       setCells(cells);

package/src/core/host/AGENTS.md

@@ -1,6 +1,16 @@
 # Host core
 
-Production code may import host modules, `jotai/vanilla`, `zod`, and
-`@/utils/Logger`. Enforced by `marimo/host-imports`. Tests are exempt.
+Production code may import host modules, `@marimo-team/client`,
+`jotai/vanilla`, `zod`, and `@/utils/Logger`. Enforced by
+`marimo/host-imports`. Tests are exempt.
+
+The protocol schemas and types live in `@marimo-team/client/protocol`.
 
 Define interfaces here and inject application behavior through adapters.
+
+`notebook-client.ts` follows the notebook with `NotebookEvents` from
+`@marimo-team/client`, over a `NotebookTransport` the editor supplies. The
+editor's transports for its session routes live in
+`core/websocket/transports/notebook.ts`. Each kind of object is a `NotebookResource`
+in `resources/` with pure handlers keyed by event name, composed with
+`createNotebookReducer`. The client holds one state atom and exposes derived atoms.

package/src/core/host/__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/host/__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");
+  });
+});

package/src/core/host/__tests__/fixtures.ts

@@ -0,0 +1,41 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import NodeFs from "node:fs";
+import NodeUrl from "node:url";
+import * as Protocol from "@marimo-team/client/protocol";
+
+/** Read the protocol's portable examples as transport messages. */
+export function fixtureMessages(name: string) {
+  const path = new NodeUrl.URL(
+    `../../../../../protocols/host/fixtures/${name}`,
+    import.meta.url,
+  );
+  return NodeFs.readFileSync(path, "utf8")
+    .split(/\r?\n\r?\n/)
+    .flatMap((frame) => {
+      const event = /^event: (.+)$/m.exec(frame)?.[1];
+      if (!event) {
+        return [];
+      }
+      const data = frame
+        .split(/\r?\n/)
+        .filter((line) => line.startsWith("data:"))
+        .map((line) => line.slice(5).trimStart())
+        .join("\n");
+      const payload: unknown = JSON.parse(data);
+      return [
+        {
+          event,
+          id: /^id: (.+)$/m.exec(frame)?.[1],
+          data: payload,
+        },
+      ];
+    });
+}
+
+/** Read the protocol's portable examples as decoded protocol events. */
+export function fixtureEvents(name: string): Protocol.EventOutput[] {
+  return fixtureMessages(name).map(({ id, ...message }) =>
+    Protocol.Event.parse(id === undefined ? message : { ...message, id }),
+  );
+}