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

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

@@ -0,0 +1,139 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Result } from "@marimo-team/client";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, type createStore } from "jotai/vanilla";
+import { vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { invariant } from "@/utils/invariant";
+import type { Document } from "../../host/resources/document";
+import type { DocumentHost } from "../model";
+import { createTestHost, type EditRequest } from "./test-host";
+
+export interface PendingRequest extends EditRequest {
+  readonly signal: AbortSignal;
+  process(): void;
+  deliverResponse(): void;
+  loseResponse(): void;
+}
+
+export function createTestNotebook(
+  store: ReturnType<typeof createStore>,
+  initial: Document,
+) {
+  const host = createTestHost(initial);
+  const source = atom<Document | null>(initial);
+  const requests: PendingRequest[] = [];
+  const waiting: Deferred<PendingRequest>[] = [];
+  const documents: Document[] = [];
+  function apply(changes: readonly Protocol.Change[]) {
+    const { result, update } = host.apply(changes);
+    invariant(result.ok, "Another writer's edit was rejected");
+    if (update !== undefined) {
+      documents.push(update);
+    }
+  }
+  const edit = vi.fn<DocumentHost["edit"]>((notebookId, changes, options) => {
+    options.signal.throwIfAborted();
+    const reply = new Deferred<Result<Protocol.EditResponseOutput>>();
+    let outcome: Result<Protocol.EditResponseOutput> | undefined;
+    const abort = () => reply.reject(options.signal.reason);
+    const finish = () => options.signal.removeEventListener("abort", abort);
+    options.signal.addEventListener("abort", abort, { once: true });
+    void reply.promise.then(finish, finish);
+    const request: PendingRequest = {
+      notebookId,
+      changes: structuredClone(changes),
+      idempotencyKey: options.idempotencyKey,
+      signal: options.signal,
+      process() {
+        invariant(outcome === undefined, "Request was already processed");
+        // Cancellation of the client's wait does not roll back a server commit.
+        const processed = host.process(request);
+        outcome = processed.result;
+        if (processed.update !== undefined) {
+          documents.push(processed.update);
+        }
+      },
+      deliverResponse() {
+        invariant(
+          outcome !== undefined,
+          "Process the request before delivering its response",
+        );
+        reply.resolve(outcome);
+      },
+      loseResponse() {
+        reply.reject(new Error("Connection lost"));
+      },
+    };
+    const next = waiting.shift();
+    if (next === undefined) {
+      requests.push(request);
+    } else {
+      next.resolve(request);
+    }
+    return reply.promise;
+  });
+  const notebook = { documentAtom: source, edit } satisfies DocumentHost;
+  return {
+    notebook,
+    control: {
+      document: host.document,
+      nextRequest(): Promise<PendingRequest> {
+        const next = requests.shift();
+        if (next !== undefined) {
+          return Promise.resolve(next);
+        }
+        const deferred = new Deferred<PendingRequest>();
+        waiting.push(deferred);
+        return deferred.promise;
+      },
+      changeCell(
+        id: Protocol.CellId,
+        patch: Partial<Omit<Protocol.CellInput, "id">>,
+      ) {
+        const cell = host.document().cells.get(id);
+        invariant(
+          cell !== undefined && cell.code !== null,
+          `No host source for cell ${id}`,
+        );
+        apply([
+          {
+            kind: "put_cell",
+            base_etag: cell.etag,
+            cell: {
+              ...cell,
+              code: cell.code,
+              ...patch,
+              config: { ...cell.config, ...patch.config },
+            },
+          },
+        ]);
+      },
+      createCell(cell: Protocol.CellInput) {
+        apply([{ kind: "put_cell", base_etag: null, cell }]);
+      },
+      deleteCell(id: Protocol.CellId) {
+        const cell = host.document().cells.get(id);
+        invariant(cell !== undefined, `No host cell ${id}`);
+        apply([{ kind: "delete_cell", id, base_etag: cell.etag }]);
+      },
+      deliverDocument() {
+        const next = documents.shift();
+        invariant(
+          next !== undefined,
+          "No document update is waiting for delivery",
+        );
+        store.set(source, next);
+      },
+      reset() {
+        documents.length = 0;
+        store.set(source, null);
+      },
+      deliverSnapshot() {
+        documents.length = 0;
+        store.set(source, host.document());
+      },
+    },
+  };
+}

package/src/core/document/__tests__/user.ts

@@ -0,0 +1,61 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { createStore } from "jotai/vanilla";
+import type { DocumentModel } from "../model";
+
+/** Translate user actions and observations without keeping any expected state. */
+export function createUser(
+  store: ReturnType<typeof createStore>,
+  model: DocumentModel,
+) {
+  function observe(promise: Promise<Protocol.Document>) {
+    // A failed assertion may end a scenario before it awaits sync; teardown
+    // should not obscure that failure with an unhandled abort rejection.
+    void promise.catch(() => undefined);
+    return promise;
+  }
+  return {
+    changeCode(id: Protocol.CellId, code: string) {
+      store.set(model.editAtom, (tx) => tx.setCode(id, code));
+    },
+    changeCells(
+      changes: Record<Protocol.CellId, Partial<Omit<Protocol.CellInput, "id">>>,
+    ) {
+      store.set(model.editAtom, (tx) => {
+        for (const [id, patch] of Object.entries(changes)) {
+          if (patch.code !== undefined) {
+            tx.setCode(id, patch.code);
+          }
+          if (patch.name !== undefined) {
+            tx.setName(id, patch.name);
+          }
+          if (patch.config !== undefined) {
+            tx.updateConfig(id, patch.config);
+          }
+        }
+      });
+    },
+    createCell(cell: Protocol.CellInput) {
+      store.set(model.editAtom, (tx) => tx.createCell(cell));
+    },
+    deleteCell(id: Protocol.CellId) {
+      store.set(model.editAtom, (tx) => tx.deleteCell(id));
+    },
+    cell(id: Protocol.CellId) {
+      const cell = store.get(model.documentAtom)?.cells.get(id);
+      return cell === undefined
+        ? undefined
+        : {
+            id: cell.id,
+            code: cell.code,
+            name: cell.name,
+            config: cell.config,
+          };
+    },
+    order: () => store.get(model.documentAtom)?.order ?? null,
+    status: () => store.get(model.syncStatusAtom),
+    sync: () => observe(model.flush()),
+    retry: () => observe(model.retry()),
+  };
+}

package/src/core/document/index.ts

@@ -6,5 +6,7 @@ export type {
   DocumentEditor,
   DocumentHost,
   DocumentModel,
-  DocumentSyncState,
+  DocumentSyncStatus,
+  EditableCell,
+  EditableDocument,
 } from "./model";

package/src/core/document/model.ts

@@ -6,18 +6,19 @@ import { createDocumentState, type DocumentState } from "./state";
 import {
   createDocumentSync,
   type DocumentHost,
-  type DocumentSyncState,
+  type DocumentSyncStatus,
 } from "./sync";
 
 export type { DocumentEdit, DocumentEditor } from "./state";
-export type { DocumentHost, DocumentSyncState } from "./sync";
+export type { EditableCell, EditableDocument } from "./draft";
+export type { DocumentHost, DocumentSyncStatus } from "./sync";
 
 /** The editable document and its synchronization with the host. */
 export interface DocumentModel extends Pick<
   DocumentState,
-  "cellAtom" | "orderAtom" | "appConfigAtom" | "editAtom"
+  "documentAtom" | "editAtom"
 > {
-  readonly syncStateAtom: Atom<DocumentSyncState>;
+  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. */
@@ -42,11 +43,9 @@ export function createDocumentModel({
   });
   const sync = createDocumentSync({ store, notebook, state, signal });
   return {
-    cellAtom: state.cellAtom,
-    orderAtom: state.orderAtom,
-    appConfigAtom: state.appConfigAtom,
+    documentAtom: state.documentAtom,
     editAtom: state.editAtom,
-    syncStateAtom: sync.stateAtom,
+    syncStatusAtom: sync.statusAtom,
     flush: sync.flush,
     retry: sync.retry,
   };

package/src/core/document/select.ts

@@ -0,0 +1,32 @@
+/* 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);
+}

package/src/core/document/state.ts

@@ -1,13 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { dequal } from "dequal";
 import {
   atom,
   type Atom,
   type createStore,
   type WritableAtom,
 } from "jotai/vanilla";
-import { selectAtom } from "jotai/vanilla/utils";
 import type { Document } from "../host/resources/document";
 import * as Protocol from "@marimo-team/client/protocol";
 import {
@@ -16,14 +14,12 @@ import {
   createCell,
   deleteCell,
   type Draft,
-  type EditableCell,
   type EditableDocument,
   editCell,
   EMPTY_DRAFT,
   isEmptyDraft,
   type Placement,
   prune,
-  sameCell,
   setOrder,
   toChanges,
   viewDocument,
@@ -70,12 +66,8 @@ interface ModelState {
  * Switching notebook identity invalidates it, retaining each notebook's draft.
  */
 export interface DocumentState {
-  /** Hold onto the selector; its value is stable until the cell changes. */
-  cellAtom(id: Protocol.CellId): Atom<EditableCell | undefined>;
-  readonly orderAtom: Atom<readonly Protocol.CellId[] | null>;
-  readonly appConfigAtom: Atom<Readonly<
-    Protocol.Document["app_config"]
-  > | null>;
+  /** 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. */
@@ -158,23 +150,6 @@ export function createDocumentState({
       ? null
       : viewDocument(state.confirmed, draftOf(state));
   });
-  const cellAtom = (id: Protocol.CellId) =>
-    selectAtom(
-      documentAtom,
-      (document) => document?.cells.get(id),
-      (a, b) =>
-        a === b || (a !== undefined && b !== undefined && sameCell(a, b)),
-    );
-  const orderAtom = selectAtom(
-    documentAtom,
-    (document) => document?.order ?? null,
-    dequal,
-  );
-  const appConfigAtom = selectAtom(
-    stateAtom,
-    (state) => state.confirmed?.revision.app_config ?? null,
-    dequal,
-  );
   const changesAtom = atom((get): readonly Protocol.Change[] | null => {
     const state = get(stateAtom);
     return state.confirmed === null
@@ -304,9 +279,7 @@ export function createDocumentState({
   }
 
   return {
-    cellAtom,
-    orderAtom,
-    appConfigAtom,
+    documentAtom,
     changesAtom,
     pendingAtom,
     editAtom,

package/src/core/document/sync.ts

@@ -9,7 +9,7 @@ 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 =
+export type DocumentSyncStatus =
   | {
       readonly status:
         | "idle"
@@ -26,7 +26,7 @@ export type DocumentSyncState =
     };
 
 interface Submission {
-  readonly edit: DocumentSubmission;
+  readonly submission: DocumentSubmission;
   readonly idempotencyKey: string;
   readonly startedAt: number;
   readonly controller: AbortController;
@@ -85,7 +85,7 @@ export function createDocumentSync({
 }) {
   signal.throwIfAborted();
   const source = notebook.documentAtom;
-  const stateAtom = atom<DocumentSyncState>({ status: "idle" });
+  const statusAtom = atom<DocumentSyncStatus>({ status: "idle" });
   const observed = atom((get) => ({
     document: get(source),
     pending: get(state.pendingAtom),
@@ -93,7 +93,7 @@ export function createDocumentSync({
   }));
   let active: Submission | null = null;
   let running: Promise<Protocol.Document> | null = null;
-  let failure: Extract<DocumentSyncState, { status: "error" }> | null = null;
+  let failure: Extract<DocumentSyncStatus, { status: "error" }> | null = null;
   let timer: ReturnType<typeof setTimeout> | undefined;
   let refreshing = false;
 
@@ -118,7 +118,7 @@ export function createDocumentSync({
       if (
         active &&
         current &&
-        current.revision.notebook_id !== active.edit.notebookId
+        current.revision.notebook_id !== active.submission.notebookId
       ) {
         active.controller.abort(
           new Error("The notebook changed during synchronization"),
@@ -130,13 +130,13 @@ export function createDocumentSync({
       clearTimer();
       const changes = store.get(state.changesAtom);
       if (changes === null) {
-        store.set(stateAtom, { status: "unavailable" });
+        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(stateAtom, { status: "idle" });
+        store.set(statusAtom, { status: "idle" });
       } else {
-        store.set(stateAtom, { status: "pending" });
+        store.set(statusAtom, { status: "pending" });
         timer = setTimeout(() => {
           void flush().catch(() => undefined);
         }, DEBOUNCE_MS);
@@ -155,7 +155,7 @@ export function createDocumentSync({
           "The edit's retry window expired; its outcome is unknown",
         );
       }
-      store.set(stateAtom, { status: "sending" });
+      store.set(statusAtom, { status: "sending" });
       const controller = new AbortController();
       const abort = () => controller.abort(request.controller.signal.reason);
       request.controller.signal.addEventListener("abort", abort, {
@@ -170,10 +170,14 @@ export function createDocumentSync({
         signal.throwIfAborted();
         request.controller.signal.throwIfAborted();
         result = await abortable(
-          notebook.edit(request.edit.notebookId, request.edit.changes, {
-            idempotencyKey: request.idempotencyKey,
-            signal: controller.signal,
-          }),
+          notebook.edit(
+            request.submission.notebookId,
+            request.submission.changes,
+            {
+              idempotencyKey: request.idempotencyKey,
+              signal: controller.signal,
+            },
+          ),
           controller.signal,
         );
       } catch (error) {
@@ -198,17 +202,17 @@ export function createDocumentSync({
           cause: result.problem,
         });
       }
-      if (result.data.document.notebook_id !== request.edit.notebookId) {
+      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.edit, request.response);
+      store.set(state.settleAtom, request.submission, request.response);
       return;
     }
   }
 
   function confirm(request: Submission): Promise<void> {
-    store.set(stateAtom, { status: "confirming" });
+    store.set(statusAtom, { status: "confirming" });
     return new Promise((resolve, reject) => {
       const check = () => {
         if (signal.aborted || request.controller.signal.aborted) {
@@ -216,7 +220,7 @@ export function createDocumentSync({
           reject(
             signal.aborted ? signal.reason : request.controller.signal.reason,
           );
-        } else if (store.get(state.pendingAtom) !== request.edit) {
+        } else if (store.get(state.pendingAtom) !== request.submission) {
           cleanup();
           resolve();
         }
@@ -247,12 +251,12 @@ export function createDocumentSync({
           if (current === null || store.get(state.changesAtom) === null) {
             throw unavailable();
           }
-          const edit = store.set(state.submitAtom);
-          if (edit === null) {
+          const submission = store.set(state.submitAtom);
+          if (submission === null) {
             return current.revision;
           }
           active = {
-            edit,
+            submission,
             idempotencyKey: idempotencyKey(),
             startedAt: Date.now(),
             controller: new AbortController(),
@@ -282,7 +286,7 @@ export function createDocumentSync({
             Date.now() - active.startedAt < RETRY_WINDOW_MS),
       };
       if (!signal.aborted) {
-        store.set(stateAtom, failure);
+        store.set(statusAtom, failure);
       }
       throw error;
     } finally {
@@ -305,7 +309,7 @@ export function createDocumentSync({
       return running;
     }
     const notebookId =
-      active?.edit.notebookId ?? store.get(source)?.revision.notebook_id;
+      active?.submission.notebookId ?? store.get(source)?.revision.notebook_id;
     if (notebookId === undefined) {
       return Promise.reject(unavailable());
     }
@@ -329,14 +333,14 @@ export function createDocumentSync({
       const rejectedAt = active.problem.edit?.document.version;
       if (
         !current ||
-        current.revision.notebook_id !== active.edit.notebookId ||
+        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.edit, null);
+      store.set(state.settleAtom, active.submission, null);
       active = null;
     }
     failure = null;
@@ -355,5 +359,5 @@ export function createDocumentSync({
     stop();
   }
   refresh();
-  return { stateAtom, flush, retry };
+  return { statusAtom, flush, retry };
 }