@marimo-team/frontend 0.25.3-dev12 → 0.25.3-dev121

package/src/core/host/notebook-client.ts

@@ -1,161 +1,260 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { Logger } from "@/utils/Logger";
+import {
+  createNotebookReducer,
+  type NotebookEvent,
+  NotebookEvents,
+  type NotebookTransport,
+  type Result,
+  StreamStopped,
+} from "@marimo-team/client";
+import type {
+  Attachment,
+  Change,
+  EditResponseOutput,
+  NotebookId,
+  NotebookOutput,
+} from "@marimo-team/client/protocol";
 import { type Atom, atom, type createStore } from "jotai/vanilla";
+import { z } from "zod";
 import {
-  type ConnectionStatus,
-  ConnectionState,
   ConnectionClosedReason,
-  type CloseDecision,
-  type NotebookTransport,
-  TransportReadyState,
+  ConnectionState,
+  type ConnectionStatus,
 } from "./connection";
-import { parseNotebookEvent } from "./events";
+import { attachments } from "./resources/attachments";
+import { type Document, document } from "./resources/document";
+
+/** The kinds the editor reads, with marimo's own notifications. */
+const TYPES = ["notebook", "cell", "document", "attachment", "marimo"] as const;
+
+type EditorEvent = NotebookEvent<(typeof TYPES)[number]>;
+
+/** A native marimo notification, passed through without interpretation. */
+const MarimoMessage = z.object({ marimo: z.unknown() });
 
 /** A connection to a notebook on a marimo host. */
 export interface NotebookClient {
-  /** Transport state, independent of kernel readiness. */
+  /** Client connection status, independent of application session readiness. */
   readonly connectionAtom: Atom<ConnectionStatus>;
+  /** Identity and capabilities, available while the stream is connected and ready. */
+  readonly notebookAtom: Atom<NotebookOutput | null>;
+  /**
+   * The most recently confirmed document.
+   *
+   * Null until a snapshot is complete, after reset or removal, and on
+   * connections without host document events. A temporary disconnection
+   * retains the last document. Partial updates are not exposed.
+   */
+  readonly documentAtom: Atom<Document | null>;
+  /** Apply an edit to the notebook captured by the caller, even across a reset. */
+  edit(
+    notebookId: NotebookId,
+    changes: readonly Change[],
+    options: { readonly idempotencyKey: string; readonly signal: AbortSignal },
+  ): Promise<Result<EditResponseOutput>>;
+  /** Live attachments, empty until the stream is ready or while disconnected. */
+  readonly attachmentsAtom: Atom<readonly Attachment[]>;
   /** Attempt to restore a closed connection after checking host health. */
   reconnect(): Promise<void>;
 }
 
-/** The transport and application callbacks needed to connect a notebook. */
+/** The transport and application callbacks needed to follow a notebook. */
 export interface NotebookConnection {
-  /** Create a transport that reads the URL again on each connection attempt. */
-  createTransport: (url: () => string) => NotebookTransport;
-  /** The current stream URL, including authentication and subscription options. */
-  url: () => string;
-  /** Host events require an intact stream; legacy notifications are independent. */
-  messageFormat: "host" | "session";
-  /** Abort to close the transport and clear the client's state. */
+  /** Opens one connection to the notebook's stream. */
+  transport: NotebookTransport;
+  /** Abort to close the stream and clear the client's state. */
   signal: AbortSignal;
-  /** Wait until the host can accept the initial connection. */
-  waitUntilReady: () => Promise<void>;
   /** Check whether the host is reachable before an explicit reconnect. */
   checkHealth: () => Promise<boolean>;
-  /** Decode wire text without interpreting its event type. */
-  decodeJson: (message: string) => unknown;
   /** Handle a native marimo notification. Its payload has not been validated. */
   onNotification: (notification: unknown) => void;
-  /** Report a malformed message without coupling the client to editor UI. */
-  onMessageError: (error: unknown) => void;
-  /** Choose the connection status and retry policy after the transport closes. */
-  classifyClose: (event: CloseEvent) => CloseDecision;
 }
 
-/** Create a client without opening a connection. */
-export function createNotebookClient(store: ReturnType<typeof createStore>) {
+const disconnected = {
+  state: ConnectionState.CLOSED,
+  code: ConnectionClosedReason.KERNEL_DISCONNECTED,
+  reason: "kernel not found",
+} satisfies ConnectionStatus;
+
+const malformed = {
+  ...disconnected,
+  reason: "Failed to decode notebook message",
+} satisfies ConnectionStatus;
+
+function isStatus(value: unknown): value is ConnectionStatus {
+  return typeof value === "object" && value !== null && "state" in value;
+}
+
+/**
+ * The status to show after the stream stopped.
+ *
+ * A transport says which status its stop means. An event the protocol
+ * rejects is malformed. Anything else means the host was not reached.
+ */
+function statusAfter(error: unknown): ConnectionStatus {
+  if (error instanceof StreamStopped && isStatus(error.cause)) {
+    return error.cause;
+  }
+  if (error instanceof StreamStopped && error.cause instanceof z.ZodError) {
+    return malformed;
+  }
+  return disconnected;
+}
+
+/**
+ * Create a client without opening a connection.
+ * The session supplies edits authenticated for the same host as its stream.
+ */
+export function createNotebookClient(
+  store: ReturnType<typeof createStore>,
+  { edit }: { edit?: NotebookClient["edit"] } = {},
+) {
+  // `useNotebookSession` copies this into the editor's shared `connectionAtom`,
+  // which gates features until MO-8051 derives availability separately.
   const connectionAtom = atom<ConnectionStatus>({
     state: ConnectionState.NOT_STARTED,
   });
+  // One atom holds the whole stream, which resource atoms derive from.
+  const reducer = createNotebookReducer({ document, attachments });
+  const streamAtom = atom(reducer.initial);
+  // A resumed stream is live, but presence waits for this connection's ready.
+  const readyAtom = atom(false);
+  const presentAtom = atom(
+    (get) =>
+      get(readyAtom) && get(connectionAtom).state === ConnectionState.OPEN,
+  );
+  const notebookAtom = atom((get) =>
+    get(presentAtom) ? get(streamAtom).notebook : null,
+  );
+  const documentAtom = atom(
+    (get) => get(streamAtom).resources.document.current,
+  );
+  const attachmentsAtom = atom((get) =>
+    get(presentAtom)
+      ? get(streamAtom).resources.attachments
+      : attachments.initial,
+  );
   let reconnect: (() => Promise<void>) | undefined;
   const client: NotebookClient = {
     connectionAtom,
+    notebookAtom,
+    documentAtom,
+    edit: async (notebookId, changes, options) => {
+      options.signal.throwIfAborted();
+      const currentId = store.get(streamAtom).notebookId;
+      if (currentId !== null && currentId !== notebookId) {
+        throw new Error("The notebook changed before the edit was sent");
+      }
+      if (edit === undefined) {
+        return {
+          ok: false,
+          status: 501,
+          problem: {
+            type: "tag:marimo.io,2026:host/unsupported-operation",
+            title: "Document editing is unavailable on this connection",
+            status: 501,
+          },
+        };
+      }
+      return edit(notebookId, changes, options);
+    },
+    attachmentsAtom,
     reconnect: async () => reconnect?.(),
   };
 
+  function setStatus(status: ConnectionStatus) {
+    store.set(readyAtom, false);
+    store.set(connectionAtom, status);
+  }
+
   function connect({
-    createTransport,
-    url,
-    messageFormat,
+    transport,
     signal,
-    waitUntilReady,
     checkHealth,
     onNotification,
-    onMessageError,
-    decodeJson,
-    classifyClose,
   }: NotebookConnection) {
     if (signal.aborted) {
       return;
     }
-    let shouldRetry = true;
-    let terminal = false;
+    let events: NotebookEvents<(typeof TYPES)[number]> | undefined;
     let checkingHealth = false;
-    const transport = createTransport(url);
-    const setStatus = (status: ConnectionStatus) =>
-      store.set(connectionAtom, status);
-    const disconnected = {
-      state: ConnectionState.CLOSED,
-      code: ConnectionClosedReason.KERNEL_DISCONNECTED,
-      reason: "kernel not found",
-    } satisfies ConnectionStatus;
-
-    function retry(code?: number, reason?: string) {
-      if (!signal.aborted && !terminal && shouldRetry) {
-        shouldRetry = false;
-        transport.reconnect(code, reason);
-      }
-    }
+    // Whether the replica still holds the state the last cursor describes.
+    let resumable = true;
 
-    function onOpen() {
-      if (signal.aborted || terminal) {
-        return;
-      }
-      shouldRetry = true;
-      setStatus({ state: ConnectionState.OPEN });
-    }
-
-    function onClose(event: CloseEvent) {
-      if (signal.aborted || terminal) {
-        return;
-      }
-      Logger.warn("Notebook connection closed", event.code, event.reason);
-      const decision = classifyClose(event);
-      setStatus(decision.status);
-      if (decision.kind === "terminal") {
-        terminal = true;
-        transport.close();
-      } else if (decision.kind === "retry") {
-        retry(event.code, event.reason);
-      } else {
-        terminal = true;
+    /** Mirror one stream into the atoms until it closes. */
+    function follow(cursor?: string) {
+      if (cursor === undefined) {
+        store.set(streamAtom, reducer.initial);
       }
+      resumable = true;
+      const current = new NotebookEvents(transport, {
+        types: TYPES,
+        cursor,
+        signal,
+        // TODO(MO-8051): Gate editing and running on the notebook's operations
+        // and request results instead. The stream's status would then only
+        // mark the view as stale.
+        onStatus: (status) => {
+          if (status === "connecting") {
+            setStatus({ state: ConnectionState.CONNECTING });
+          } else if (status === "connected") {
+            setStatus({ state: ConnectionState.OPEN });
+          }
+        },
+      });
+      current.subscribe(
+        (event) => apply(current, event),
+        (error) => {
+          if (!signal.aborted && current === events) {
+            Logger.warn("Notebook stream stopped", error);
+            const status = statusAfter(error);
+            if (status === malformed) {
+              stopMalformed(current);
+            } else {
+              setStatus(status);
+            }
+          }
+        },
+      );
+      events = current;
     }
 
-    function onError(event: Event) {
-      if (signal.aborted || terminal) {
-        return;
-      }
-      Logger.warn("Notebook connection error", event);
-      setStatus(disconnected);
-      retry();
+    function stopMalformed(source: NotebookEvents<string>) {
+      source.close();
+      // The cursor can already include an event the replica rejected.
+      resumable = false;
+      store.set(streamAtom, reducer.initial);
+      setStatus(malformed);
     }
 
-    function onMessage(message: MessageEvent<string>) {
-      if (signal.aborted || terminal) {
+    function apply(source: NotebookEvents<string>, event: EditorEvent) {
+      if (event.event === "marimo") {
+        const message = MarimoMessage.safeParse(event.data);
+        if (message.success) {
+          onNotification(message.data.marimo);
+        } else {
+          Logger.error("Failed to decode notebook message", event);
+          stopMalformed(source);
+        }
         return;
       }
-      let event;
       try {
-        event = parseNotebookEvent(decodeJson(message.data));
+        store.set(streamAtom, reducer.reduce(store.get(streamAtom), event));
       } catch (error) {
-        Logger.error("Failed to decode notebook message", error);
-        if (messageFormat === "session") {
-          onMessageError(error);
-          return;
-        }
-        terminal = true;
-        setStatus({
-          ...disconnected,
-          reason: "Failed to decode notebook message",
-        });
-        transport.close();
+        // An event the replica cannot apply, such as another notebook's.
+        Logger.error("Failed to apply notebook event", error);
+        stopMalformed(source);
         return;
       }
-      if (event.event === "marimo") {
-        onNotification(event.data.marimo);
+      if (event.event === "ready" || event.event === "reset") {
+        store.set(readyAtom, event.event === "ready");
       }
     }
 
     reconnect = async () => {
-      if (
-        signal.aborted ||
-        checkingHealth ||
-        transport.readyState === TransportReadyState.OPEN ||
-        transport.readyState === TransportReadyState.CONNECTING
-      ) {
+      if (signal.aborted || checkingHealth || events?.status !== "closed") {
         return;
       }
       checkingHealth = true;
@@ -166,17 +265,13 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
           return;
         }
         if (!healthy) {
-          shouldRetry = false;
           setStatus(disconnected);
           return;
         }
-        terminal = false;
-        shouldRetry = true;
-        transport.reconnect();
+        follow(resumable ? (events.cursor ?? undefined) : undefined);
       } catch (error) {
         if (!signal.aborted) {
           Logger.error("Notebook health check failed", error);
-          shouldRetry = false;
           setStatus(disconnected);
         }
       } finally {
@@ -187,39 +282,15 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
     signal.addEventListener(
       "abort",
       () => {
-        transport.removeEventListener("open", onOpen);
-        transport.removeEventListener("close", onClose);
-        transport.removeEventListener("error", onError);
-        transport.removeEventListener("message", onMessage);
-        transport.close();
         reconnect = undefined;
+        store.set(streamAtom, reducer.initial);
         setStatus({ state: ConnectionState.NOT_STARTED });
       },
       { once: true },
     );
 
-    transport.addEventListener("open", onOpen);
-    transport.addEventListener("close", onClose);
-    transport.addEventListener("error", onError);
-    transport.addEventListener("message", onMessage);
-
-    if (transport.readyState === TransportReadyState.CLOSED) {
-      setStatus({ state: ConnectionState.CONNECTING });
-      void waitUntilReady()
-        .then(() => {
-          if (!signal.aborted) {
-            transport.reconnect();
-          }
-        })
-        .catch((error) => {
-          if (!signal.aborted) {
-            Logger.error("Healthy connection never made", error);
-            terminal = true;
-            setStatus(disconnected);
-            transport.close();
-          }
-        });
-    }
+    setStatus({ state: ConnectionState.CONNECTING });
+    follow();
   }
 
   return { client, connect };

package/src/core/host/resources/attachments.ts

@@ -0,0 +1,22 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Attachment } from "@marimo-team/client/protocol";
+import type { NotebookResource } from "@marimo-team/client";
+
+/** Attachment records retained with the stream's resume cursor. */
+export const attachments: NotebookResource<
+  readonly Attachment[],
+  "attachment"
+> = {
+  initial: [],
+  on: {
+    attachment: (state, { data: { attachment } }) =>
+      state.some((item) => item.id === attachment.id)
+        ? state.map((item) => (item.id === attachment.id ? attachment : item))
+        : [...state, attachment],
+    "attachment.removed": (state, { data }) =>
+      state.some((item) => item.id === data.id)
+        ? state.filter((item) => item.id !== data.id)
+        : state,
+  },
+};

package/src/core/host/resources/document.ts

@@ -0,0 +1,101 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { NotebookResource } from "@marimo-team/client";
+
+/**
+ * The notebook's content at one revision.
+ *
+ * The revision record carries the cell order, app config, etags, and
+ * version. Each cell in its order has a record in `cells`.
+ */
+export interface Document {
+  readonly revision: Protocol.Document;
+  readonly cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>;
+}
+
+/** The records of a revision still being received. */
+interface PendingDocument {
+  /** Null until the `document` event that closes the group arrives. */
+  readonly revision: Protocol.Document | null;
+  readonly cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>;
+}
+
+/**
+ * The confirmed document and the revision being assembled after it.
+ *
+ * A snapshot's records are pending until `ready`. A live revision is pending
+ * from its first cell until its `document` event.
+ */
+interface DocumentState {
+  readonly current: Document | null;
+  readonly pending: PendingDocument | null;
+}
+
+function pendingOf(state: DocumentState): PendingDocument {
+  return (
+    state.pending ?? {
+      revision: state.current?.revision ?? null,
+      cells: state.current?.cells ?? new Map(),
+    }
+  );
+}
+
+/** Confirm a revision whose records are complete. */
+function publish(
+  revision: Protocol.Document,
+  cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>,
+): DocumentState {
+  const order = new Set(revision.order);
+  if (
+    order.size !== revision.order.length ||
+    order.size !== cells.size ||
+    [...order].some((id) => !cells.has(id)) ||
+    [...cells.values()].some((cell) => cell.version > revision.version)
+  ) {
+    throw new Error("Document records do not match its revision and order");
+  }
+  return { current: { revision, cells }, pending: null };
+}
+
+export const document: NotebookResource<DocumentState, "cell" | "document"> = {
+  initial: { current: null, pending: null },
+  on: {
+    cell: (state, { data }) => {
+      const pending = pendingOf(state);
+      return {
+        ...state,
+        pending: {
+          revision: pending.revision,
+          cells: new Map(pending.cells).set(data.cell.id, data.cell),
+        },
+      };
+    },
+    "cell.removed": (state, { data }) => {
+      const pending = pendingOf(state);
+      if (!pending.cells.has(data.id)) {
+        return state;
+      }
+      const cells = new Map(pending.cells);
+      cells.delete(data.id);
+      return { ...state, pending: { ...pending, cells } };
+    },
+    document: (state, { data }, { phase }) => {
+      const { cells } = pendingOf(state);
+      return phase === "snapshot"
+        ? { ...state, pending: { revision: data.document, cells } }
+        : publish(data.document, cells);
+    },
+    ready: (state, _event, { phase }) => {
+      if (state.pending === null) {
+        return state;
+      }
+      if (phase === "live") {
+        throw new Error("Stream ended with an incomplete document update");
+      }
+      return state.pending.revision === null
+        ? { current: null, pending: null }
+        : publish(state.pending.revision, state.pending.cells);
+    },
+  },
+};

package/src/core/host/static.ts

@@ -0,0 +1,70 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { NotebookBackend } from "./backend";
+
+/** A document snapshot with its own identity and revision metadata. */
+export interface StaticNotebookSnapshot {
+  readonly notebook: Omit<Protocol.Notebook, "operations">;
+  readonly document: Protocol.Document;
+  readonly cells: readonly Protocol.Cell[];
+}
+
+/**
+ * Serve a document without a server, runtime, or mutation operations.
+ * Callers and readers must treat the snapshot records as immutable.
+ */
+export function createStaticNotebookBackend(
+  snapshot: StaticNotebookSnapshot,
+): NotebookBackend {
+  const frames: Protocol.Event[] = [
+    {
+      event: "notebook",
+      data: { notebook: { ...snapshot.notebook, operations: [] } },
+    },
+    ...snapshot.cells.map((cell): Protocol.CellEvent => ({
+      event: "cell",
+      data: { cell },
+    })),
+    { event: "document", data: { document: snapshot.document } },
+    { event: "ready", id: "static", data: {} },
+  ];
+
+  return {
+    transport: ({ cursor, signal }, { open, message }) =>
+      new Promise((resolve) => {
+        if (signal.aborted) {
+          resolve();
+          return;
+        }
+        // Finishing would ask NotebookEvents to reconnect. A snapshot stays open
+        // until its reader closes, even though there are no more updates.
+        signal.addEventListener("abort", () => resolve(), { once: true });
+        open();
+        if (cursor !== null && !signal.aborted) {
+          message({ event: "reset", data: {} });
+        }
+        for (const frame of frames) {
+          if (signal.aborted) {
+            return;
+          }
+          message(frame);
+        }
+      }),
+    client: {
+      edit: async (_changes, options) => {
+        options?.signal?.throwIfAborted();
+        return {
+          ok: false,
+          status: 501,
+          problem: {
+            type: "tag:marimo.io,2026:host/unsupported-operation",
+            title: "Unsupported operation",
+            status: 501,
+            detail: "Static notebooks do not support editing",
+          },
+        };
+      },
+    },
+  };
+}