@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev284

package/src/core/notebook/document/hooks.tsx

@@ -0,0 +1,118 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { CellId } from "@marimo-team/client/protocol";
+import { useAtomValue, useSetAtom, useStore } from "jotai";
+import { createContext, type ReactNode, useContext, useMemo } from "react";
+import { cellCodeBinding } from "./codemirror";
+import type { EditableCell } from "./draft";
+import type { DocumentModel } from "./model";
+import { selectCellField, selectOrder } from "./select";
+
+const DocumentContext = createContext<DocumentModel | null>(null);
+
+/**
+ * Expose a session-owned model to React. Construct the model with the same store
+ * as the enclosing Jotai Provider. The session owner starts it; mounting or
+ * unmounting this provider does not change the model's lifetime.
+ *
+ * ```tsx
+ * <Provider store={store}>
+ *   <DocumentProvider model={model}>
+ *     <Editor />
+ *   </DocumentProvider>
+ * </Provider>
+ * ```
+ */
+export function DocumentProvider({
+  model,
+  children,
+}: {
+  model: DocumentModel;
+  children: ReactNode;
+}) {
+  return (
+    <DocumentContext.Provider value={model}>
+      {children}
+    </DocumentContext.Provider>
+  );
+}
+
+function useDocumentModel() {
+  const model = useContext(DocumentContext);
+  if (model === null) {
+    throw new Error("Document hooks require a DocumentProvider");
+  }
+  return model;
+}
+
+/** Null while the document is unavailable; an empty array for an empty document. */
+export function useNotebookCellOrder() {
+  const { documentAtom } = useDocumentModel();
+  return useAtomValue(useMemo(() => selectOrder(documentAtom), [documentAtom]));
+}
+
+function useCellField<K extends keyof EditableCell>(id: CellId, field: K) {
+  const { documentAtom } = useDocumentModel();
+  return useAtomValue(
+    useMemo(
+      () => selectCellField(documentAtom, id, field),
+      [documentAtom, id, field],
+    ),
+  );
+}
+
+/** Undefined for an absent cell; null when the host withholds its source. */
+export function useNotebookCellCode(id: CellId) {
+  return useCellField(id, "code");
+}
+
+/** Undefined when the cell or document is unavailable. */
+export function useNotebookCellName(id: CellId) {
+  return useCellField(id, "name");
+}
+
+/** Undefined when the cell or document is unavailable. */
+export function useNotebookCellConfig(id: CellId) {
+  return useCellField(id, "config");
+}
+
+/**
+ * A CodeMirror extension that binds the editor's text to a cell's code.
+ * Stable until the cell id or document model changes.
+ *
+ * ```tsx
+ * const extension = useNotebookCellCodeExtension(cellId);
+ * ```
+ */
+export function useNotebookCellCodeExtension(id: CellId) {
+  const model = useDocumentModel();
+  const store = useStore();
+  return useMemo(
+    () => cellCodeBinding({ store, model, cellId: id }),
+    [model, store, id],
+  );
+}
+
+/**
+ * A stable edit function. Synchronous callbacks publish atomically or roll back.
+ *
+ * ```ts
+ * const edit = useNotebookDocumentEdit();
+ * edit((tx) => tx.setCode(cellId, code));
+ * ```
+ */
+export function useNotebookDocumentEdit() {
+  return useSetAtom(useDocumentModel().editAtom);
+}
+
+/**
+ * Sync status and explicit recovery. flush waits for edits, including newer
+ * typing, to reach the confirmed source and returns the revision for Run/Save.
+ * Rejections retain local edits; retry explicitly resubmits them. Requests
+ * with an unknown outcome reuse their original key within the retry window.
+ */
+export function useNotebookDocumentSync() {
+  const model = useDocumentModel();
+  const status = useAtomValue(model.syncStatusAtom);
+  return { status, flush: model.flush, retry: model.retry };
+}

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

@@ -0,0 +1,24 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The editable document, its drafts, and its synchronization with the host. */
+
+export {
+  DocumentProvider,
+  useNotebookCellCode,
+  useNotebookCellCodeExtension,
+  useNotebookCellConfig,
+  useNotebookCellName,
+  useNotebookCellOrder,
+  useNotebookDocumentEdit,
+  useNotebookDocumentSync,
+} from "./hooks";
+export {
+  createDocumentModel,
+  type DocumentEdit,
+  type DocumentEditor,
+  type DocumentModel,
+  type DocumentSyncStatus,
+  type EditableCell,
+  type EditableDocument,
+} from "./model";
+export type { Placement } from "./draft";

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

@@ -0,0 +1,102 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import {
+  type Atom,
+  atom,
+  type createStore,
+  type WritableAtom,
+} from "jotai/vanilla";
+import type { Document, NotebookBackend } from "../replica";
+import type { EditableDocument } from "./draft";
+import {
+  createDocumentState,
+  type DocumentEdit,
+  type DocumentState,
+} from "./state";
+import { createDocumentSync, type DocumentSyncStatus } from "./sync";
+
+export type { DocumentEdit, DocumentEditor } from "./state";
+export type { EditableCell, EditableDocument } from "./draft";
+export type { DocumentSyncStatus } from "./sync";
+
+/** The editable document and its synchronization with the host. */
+export interface DocumentModel {
+  /** Confirmed host state with local drafts applied, null while unavailable or stopped. */
+  readonly documentAtom: Atom<EditableDocument | null>;
+  /** Ignored while stopped. */
+  readonly editAtom: WritableAtom<null, [edit: DocumentEdit], void>;
+  readonly syncStatusAtom: Atom<DocumentSyncStatus>;
+  /** Drain local edits and return the confirmed revision. */
+  flush(): Promise<Protocol.Document>;
+  /** Retry a rejection against fresh state, or repeat an uncertain request exactly. */
+  retry(): Promise<Protocol.Document>;
+  /**
+   * Edit and synchronize until the signal aborts, which discards local drafts.
+   * The model can start again once stopped.
+   */
+  start(options: { signal: AbortSignal }): void;
+}
+
+interface Running {
+  readonly state: DocumentState;
+  readonly sync: ReturnType<typeof createDocumentSync>;
+}
+
+/** Create without subscribing to anything; editing begins on start. */
+export function createDocumentModel({
+  store,
+  documentAtom: source,
+  backendAtom,
+}: {
+  store: ReturnType<typeof createStore>;
+  /** The document the host last confirmed. */
+  documentAtom: Atom<Document | null>;
+  /** The backend that confirmed it, whose commands reach its host. */
+  backendAtom: Atom<NotebookBackend | null>;
+}): DocumentModel {
+  const runningAtom = atom<Running | null>(null);
+
+  function stopped(): Promise<never> {
+    return Promise.reject(new Error("The document model is not started"));
+  }
+
+  return {
+    documentAtom: atom((get) => {
+      const running = get(runningAtom);
+      return running ? get(running.state.documentAtom) : null;
+    }),
+    editAtom: atom(null, (get, set, edit: DocumentEdit) => {
+      const running = get(runningAtom);
+      if (running) {
+        set(running.state.editAtom, edit);
+      }
+    }),
+    syncStatusAtom: atom((get) => {
+      const running = get(runningAtom);
+      return running ? get(running.sync.statusAtom) : { status: "unavailable" };
+    }),
+    flush: () => store.get(runningAtom)?.sync.flush() ?? stopped(),
+    retry: () => store.get(runningAtom)?.sync.retry() ?? stopped(),
+    start({ signal }) {
+      signal.throwIfAborted();
+      if (store.get(runningAtom)) {
+        throw new Error("The document model is already started");
+      }
+      const state = createDocumentState({ store, source, signal });
+      const sync = createDocumentSync({
+        store,
+        documentAtom: source,
+        backendAtom,
+        state,
+        signal,
+      });
+      const stop = () => store.set(runningAtom, null);
+      signal.addEventListener("abort", stop, { once: true });
+      store.set(runningAtom, { state, sync });
+      if (signal.aborted) {
+        stop();
+      }
+    },
+  };
+}

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

@@ -0,0 +1,41 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import { dequal } from "dequal";
+import type { Atom } from "jotai/vanilla";
+import { selectAtom } from "jotai/vanilla/utils";
+import { type EditableCell, type EditableDocument, sameCell } from "./draft";
+
+/**
+ * Select one cell of a document.
+ *
+ * The atom notifies when the cell's code, name, or config changes. A new host
+ * record with the same values leaves it unchanged. Each call creates an atom,
+ * so keep it for as long as its subscription.
+ */
+export function selectCell(
+  document: Atom<EditableDocument | null>,
+  id: Protocol.CellId,
+): Atom<EditableCell | undefined> {
+  return selectAtom(
+    document,
+    (value) => value?.cells.get(id),
+    (a, b) => a === b || (a !== undefined && b !== undefined && sameCell(a, b)),
+  );
+}
+
+/** Select a document's cell order, null while the document is unavailable. */
+export function selectOrder(
+  document: Atom<EditableDocument | null>,
+): Atom<readonly Protocol.CellId[] | null> {
+  return selectAtom(document, (value) => value?.order ?? null, dequal);
+}
+
+/** Select one field of a cell, unchanged until that field's value changes. */
+export function selectCellField<K extends keyof EditableCell>(
+  document: Atom<EditableDocument | null>,
+  id: Protocol.CellId,
+  field: K,
+): Atom<EditableCell[K] | undefined> {
+  return selectAtom(document, (value) => value?.cells.get(id)?.[field], dequal);
+}

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

@@ -0,0 +1,378 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  atom,
+  type Atom,
+  type createStore,
+  type WritableAtom,
+} from "jotai/vanilla";
+import type { Document } from "../replica";
+import * as Protocol from "@marimo-team/client/protocol";
+import {
+  acknowledge,
+  type CellEdit,
+  createCell,
+  deleteCell,
+  restoreCell,
+  type Draft,
+  type EditableDocument,
+  editCell,
+  EMPTY_DRAFT,
+  isEmptyDraft,
+  type Placement,
+  prune,
+  setOrder,
+  toChanges,
+  viewDocument,
+} from "./draft";
+
+/** Operations on a working draft, valid only during the edit callback. */
+export interface DocumentEditor {
+  setCode(id: Protocol.CellId, code: string): void;
+  setName(id: Protocol.CellId, name: string): void;
+  updateConfig(id: Protocol.CellId, patch: Partial<Protocol.CellConfig>): void;
+  /** Create a cell, or bring back one this draft deletes with this content. */
+  createCell(cell: Protocol.CellInput, placement?: Placement): void;
+  deleteCell(id: Protocol.CellId): void;
+  /** A permutation of the cells present at this point in the transaction. */
+  setOrder(order: readonly Protocol.CellId[]): void;
+}
+
+/** Synchronous; a thrown error discards the entire transaction. */
+export type DocumentEdit = (editor: DocumentEditor) => void;
+
+/** A captured edit request. Return this same object when reporting its result. */
+export interface DocumentSubmission {
+  readonly notebookId: Protocol.NotebookId;
+  readonly changes: readonly Protocol.Change[];
+}
+
+type ResponseRevision = Pick<Protocol.Document, "notebook_id" | "version">;
+
+interface PendingSubmission {
+  readonly submission: DocumentSubmission;
+  readonly draft: Draft;
+  readonly responseRevision: ResponseRevision | null;
+}
+
+interface ModelState {
+  readonly confirmed: Document | null;
+  readonly drafts: ReadonlyMap<Protocol.NotebookId, Draft>;
+  readonly pending: PendingSubmission | null;
+}
+
+/**
+ * Local edits over confirmed host state, with one submission in flight.
+ *
+ * A reset retains the submission so its receipt can reconcile a new snapshot.
+ * Switching notebook identity invalidates it, retaining each notebook's draft.
+ */
+export interface DocumentState {
+  /** Confirmed host state with local drafts applied, null while unavailable. */
+  readonly documentAtom: Atom<EditableDocument | null>;
+  /** Null when the document or source needed by a put is unavailable. */
+  readonly changesAtom: Atom<readonly Protocol.Change[] | null>;
+  /** Remains pending until both the response and its streamed revision arrive. */
+  readonly pendingAtom: Atom<DocumentSubmission | null>;
+  readonly editAtom: WritableAtom<null, [edit: DocumentEdit], void>;
+  /** Capture a request, or null when unavailable, unchanged, or already sending. */
+  readonly submitAtom: WritableAtom<null, [], DocumentSubmission | null>;
+  /**
+   * Report the response revision, or null for a rejected request.
+   *
+   * Success waits for a complete streamed document at that version or later.
+   * Duplicate replies and replies for another notebook are ignored. Retry an
+   * uncertain request with the same body and idempotency key; do not reject it
+   * or capture a new request until its outcome is known.
+   */
+  readonly settleAtom: WritableAtom<
+    null,
+    [submission: DocumentSubmission, responseRevision: ResponseRevision | null],
+    void
+  >;
+}
+
+function draftOf(state: ModelState): Draft {
+  return state.confirmed === null
+    ? EMPTY_DRAFT
+    : (state.drafts.get(state.confirmed.revision.notebook_id) ?? EMPTY_DRAFT);
+}
+
+function withDraft(state: ModelState, draft: Draft): ModelState {
+  if (state.confirmed === null || draft === draftOf(state)) {
+    return state;
+  }
+  const drafts = new Map(state.drafts);
+  const notebookId = state.confirmed.revision.notebook_id;
+  if (isEmptyDraft(draft)) {
+    drafts.delete(notebookId);
+  } else {
+    drafts.set(notebookId, draft);
+  }
+  return { ...state, drafts };
+}
+
+/** Publish the confirmed document and any released fields together. */
+function reconcile(state: ModelState): ModelState {
+  const { confirmed, pending } = state;
+  if (confirmed === null) {
+    return state;
+  }
+  let draft = draftOf(state);
+  if (
+    pending !== null &&
+    pending.responseRevision !== null &&
+    confirmed.revision.version >= pending.responseRevision.version
+  ) {
+    draft = acknowledge(draft, pending.draft);
+    state = { ...state, pending: null };
+  }
+  return withDraft(state, prune(confirmed, draft, state.pending?.draft));
+}
+
+/** Internal draft and submission state for one editing lifetime. */
+export function createDocumentState({
+  store,
+  source,
+  signal,
+}: {
+  store: ReturnType<typeof createStore>;
+  source: Atom<Document | null>;
+  signal: AbortSignal;
+}): DocumentState {
+  signal.throwIfAborted();
+  const stateAtom = atom<ModelState>({
+    confirmed: store.get(source),
+    drafts: new Map(),
+    pending: null,
+  });
+  const documentAtom = atom((get): EditableDocument | null => {
+    const state = get(stateAtom);
+    return state.confirmed === null
+      ? null
+      : viewDocument(state.confirmed, draftOf(state));
+  });
+  const changesAtom = atom((get): readonly Protocol.Change[] | null => {
+    const state = get(stateAtom);
+    return state.confirmed === null
+      ? null
+      : toChanges(state.confirmed, draftOf(state));
+  });
+  const pendingAtom = atom((get) => get(stateAtom).pending?.submission ?? null);
+
+  const unsubscribe = store.sub(source, () => {
+    const confirmed = store.get(source);
+    store.set(stateAtom, (state) =>
+      reconcile({
+        ...state,
+        confirmed,
+        pending:
+          confirmed !== null &&
+          state.pending !== null &&
+          confirmed.revision.notebook_id !== state.pending.submission.notebookId
+            ? null
+            : state.pending,
+      }),
+    );
+  });
+
+  let editing = false;
+  const editAtom = atom(null, (get, set, edit: DocumentEdit) => {
+    if (editing) {
+      throw new Error("Document transactions cannot be nested");
+    }
+    const state = get(stateAtom);
+    if (signal.aborted || state.confirmed === null) {
+      return;
+    }
+    editing = true;
+    try {
+      const draft = editDocument(
+        state.confirmed,
+        draftOf(state),
+        edit,
+        state.pending?.draft,
+      );
+      if (
+        signal.aborted ||
+        get(stateAtom) !== state ||
+        get(source) !== state.confirmed
+      ) {
+        throw new Error("The document changed during the transaction");
+      }
+      set(stateAtom, withDraft(state, draft));
+    } finally {
+      editing = false;
+    }
+  });
+
+  const submitAtom = atom(null, (get, set): DocumentSubmission | null => {
+    const state = get(stateAtom);
+    if (signal.aborted || state.confirmed === null || state.pending !== null) {
+      return null;
+    }
+    const draft = draftOf(state);
+    const changes = toChanges(state.confirmed, draft);
+    if (changes === null) {
+      return null;
+    }
+    if (changes.length === 0) {
+      set(stateAtom, withDraft(state, EMPTY_DRAFT));
+      return null;
+    }
+    const submission = {
+      notebookId: state.confirmed.revision.notebook_id,
+      changes,
+    };
+    set(stateAtom, {
+      ...state,
+      pending: { submission, draft, responseRevision: null },
+    });
+    return submission;
+  });
+
+  const settleAtom = atom(
+    null,
+    (
+      get,
+      set,
+      submission: DocumentSubmission,
+      responseRevision: ResponseRevision | null,
+    ) => {
+      const state = get(stateAtom);
+      if (
+        signal.aborted ||
+        state.pending?.submission !== submission ||
+        state.pending.responseRevision !== null
+      ) {
+        return;
+      }
+      if (
+        responseRevision !== null &&
+        responseRevision.notebook_id !== submission.notebookId
+      ) {
+        throw new Error("Edit response belongs to another notebook");
+      }
+      set(
+        stateAtom,
+        reconcile({
+          ...state,
+          pending:
+            responseRevision === null
+              ? null
+              : { ...state.pending, responseRevision },
+        }),
+      );
+    },
+  );
+
+  function stop() {
+    unsubscribe();
+    signal.removeEventListener("abort", stop);
+    store.set(stateAtom, {
+      confirmed: null,
+      drafts: new Map(),
+      pending: null,
+    });
+  }
+  signal.addEventListener("abort", stop, { once: true });
+  if (signal.aborted) {
+    stop();
+  }
+
+  return {
+    documentAtom,
+    changesAtom,
+    pendingAtom,
+    editAtom,
+    submitAtom,
+    settleAtom,
+  };
+}
+
+function editDocument(
+  confirmed: Document,
+  draft: Draft,
+  edit: DocumentEdit,
+  sent: Draft = EMPTY_DRAFT,
+): Draft {
+  let next = draft;
+  let active = true;
+
+  function assertActive() {
+    if (!active) {
+      throw new Error("The document transaction has ended");
+    }
+  }
+
+  function hasCell(id: Protocol.CellId): boolean {
+    const local = next.cells.get(id);
+    return (
+      local?.kind === "create" ||
+      (local?.kind !== "delete" && confirmed.cells.has(id))
+    );
+  }
+
+  function requireCell(id: Protocol.CellId) {
+    assertActive();
+    if (!hasCell(id)) {
+      throw new Error(`Cell ${id} does not exist`);
+    }
+  }
+
+  function updateCell(id: Protocol.CellId, patch: CellEdit) {
+    requireCell(id);
+    next = editCell(confirmed, next, id, patch);
+  }
+
+  const editor: DocumentEditor = {
+    setCode: (id, code) => updateCell(id, { code }),
+    setName: (id, name) => updateCell(id, { name }),
+    updateConfig: (id, patch) =>
+      updateCell(id, { config: Protocol.CellConfig.partial().parse(patch) }),
+    createCell: (cell, placement = {}) => {
+      assertActive();
+      if (hasCell(cell.id)) {
+        throw new Error(`Cell ${cell.id} already exists`);
+      }
+      if (placement.after !== undefined && placement.before !== undefined) {
+        throw new Error("Specify either before or after for cell placement");
+      }
+      const anchor = placement.after ?? placement.before;
+      if (anchor !== undefined) {
+        requireCell(anchor);
+      }
+      next = confirmed.cells.has(cell.id)
+        ? restoreCell(confirmed, next, cell, { ...placement }, sent)
+        : createCell(next, cell, { ...placement });
+    },
+    deleteCell: (id) => {
+      requireCell(id);
+      next = deleteCell(next, id, sent);
+    },
+    setOrder: (order) => {
+      assertActive();
+      const { cells } = viewDocument(confirmed, next);
+      if (
+        order.length !== cells.size ||
+        new Set(order).size !== cells.size ||
+        order.some((id) => !cells.has(id))
+      ) {
+        throw new Error("Cell order must contain every cell exactly once");
+      }
+      next = setOrder(next, order);
+    },
+  };
+
+  try {
+    const result: unknown = edit(editor);
+    if (result instanceof Promise) {
+      // Observe a rejected async callback without publishing its partial draft.
+      result.catch(() => undefined);
+      throw new Error("Document edit callbacks must be synchronous");
+    }
+    return next;
+  } finally {
+    active = false;
+  }
+}