@marimo-team/frontend 0.25.3-dev13 → 0.25.3-dev134

package/src/core/host/__tests__/static.test.ts

@@ -0,0 +1,169 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createNotebookReducer, NotebookEvents } from "@marimo-team/client";
+import * as Protocol from "@marimo-team/client/protocol";
+import { afterEach, assert, describe, expect, it, vi } from "vitest";
+import { document as documentResource } from "../resources/document";
+import {
+  createStaticNotebookBackend,
+  type StaticNotebookSnapshot,
+} from "../static";
+import { fixtureEvents } from "./fixtures";
+
+function snapshot(): StaticNotebookSnapshot {
+  const events = fixtureEvents("notebook-events.sse");
+  const notebook = events.find((event) => event.event === "notebook")?.data
+    .notebook;
+  const document = events.find((event) => event.event === "document")?.data
+    .document;
+  assert(notebook && document);
+  return {
+    notebook,
+    document,
+    cells: events.flatMap((event) =>
+      event.event === "cell" ? [event.data.cell] : [],
+    ),
+  };
+}
+
+const streams: NotebookEvents[] = [];
+afterEach(async () => {
+  const closed = streams.splice(0);
+  closed.forEach((stream) => stream.close());
+  await Promise.all(closed.map((stream) => stream.closed));
+  vi.useRealTimers();
+});
+
+async function read(
+  backend: ReturnType<typeof createStaticNotebookBackend>,
+  cursor?: string,
+) {
+  const stream = new NotebookEvents(backend.transport, { cursor });
+  streams.push(stream);
+  const events: Protocol.EventOutput[] = [];
+  await new Promise<void>((resolve, reject) => {
+    stream.subscribe((event) => {
+      events.push(Protocol.Event.parse(event));
+      if (event.event === "ready") {
+        resolve();
+      }
+    }, reject);
+  });
+  return { stream, events };
+}
+
+describe("static notebook backend", () => {
+  it("supplies a complete document through the normal event reader and reducer", async () => {
+    const source = snapshot();
+    const { stream, events } = await read(createStaticNotebookBackend(source));
+    const reducer = createNotebookReducer({ document: documentResource });
+    const state = events.reduce(reducer.reduce, reducer.initial);
+    expect(events.map((event) => event.event)).toEqual([
+      "notebook",
+      "cell",
+      "cell",
+      "cell",
+      "document",
+      "ready",
+    ]);
+    expect(state.notebook).toEqual({ ...source.notebook, operations: [] });
+    expect(state.resources.document.current).toEqual({
+      revision: source.document,
+      cells: new Map(source.cells.map((cell) => [cell.id, cell])),
+    });
+    expect(stream.status).toBe("connected");
+  });
+
+  it("preserves withheld source and its revision metadata", async () => {
+    const source = snapshot();
+    const cells = source.cells.map((cell) => ({ ...cell, code: null }));
+    const { events } = await read(
+      createStaticNotebookBackend({ ...source, cells }),
+    );
+    expect(
+      events
+        .filter((event) => event.event === "cell")
+        .map((event) => event.data.cell),
+    ).toEqual(cells);
+    expect(
+      events.find((event) => event.event === "document")?.data.document,
+    ).toEqual(source.document);
+  });
+
+  it("serves an empty document as available", async () => {
+    const source = snapshot();
+    const document = { ...source.document, order: [] };
+    const { events } = await read(
+      createStaticNotebookBackend({ ...source, document, cells: [] }),
+    );
+    const reducer = createNotebookReducer({ document: documentResource });
+    expect(
+      events.reduce(reducer.reduce, reducer.initial).resources.document.current,
+    ).toEqual({
+      revision: document,
+      cells: new Map(),
+    });
+  });
+
+  it("keeps independent readers open without polling or reconnecting", async () => {
+    vi.useFakeTimers();
+    const backend = createStaticNotebookBackend(snapshot());
+    const first = await read(backend);
+    const second = await read(backend);
+    first.stream.close();
+    await first.stream.closed;
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(second.events).toEqual(first.events);
+    expect(second.stream.status).toBe("connected");
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("resets a resume cursor before replaying its snapshot", async () => {
+    const { events } = await read(
+      createStaticNotebookBackend(snapshot()),
+      "previous-10",
+    );
+    expect(events[0]).toEqual({ event: "reset", data: {} });
+    expect(events.at(-1)).toEqual({ event: "ready", id: "static", data: {} });
+  });
+
+  it("advertises no mutations and returns a protocol failure for direct edits", async () => {
+    const backend = createStaticNotebookBackend(snapshot());
+    expect(await backend.notebook(snapshot().notebook.id).edit([])).toEqual({
+      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",
+      },
+    });
+  });
+
+  it("respects cancellation before opening, during delivery, and for commands", async () => {
+    const backend = createStaticNotebookBackend(snapshot());
+    const before = new AbortController();
+    before.abort();
+    const open = vi.fn();
+    const message = vi.fn();
+    await backend.transport(
+      { cursor: null, signal: before.signal },
+      { open, message },
+    );
+    expect(open).not.toHaveBeenCalled();
+    expect(message).not.toHaveBeenCalled();
+    await expect(
+      backend
+        .notebook(snapshot().notebook.id)
+        .edit([], { signal: before.signal }),
+    ).rejects.toBe(before.signal.reason);
+    const during = new AbortController();
+    message.mockImplementation(() => during.abort());
+    await backend.transport(
+      { cursor: null, signal: during.signal },
+      { open, message },
+    );
+    expect(message).toHaveBeenCalledOnce();
+  });
+});

package/src/core/host/__tests__/wasm.test.ts

@@ -0,0 +1,162 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { NotebookEvents, StreamStopped } from "@marimo-team/client";
+import { afterEach, assert, describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { createWasmNotebookBackend, type WasmNotebookBridge } from "../wasm";
+import { fixtureEvents } from "./fixtures";
+
+const streams: NotebookEvents[] = [];
+afterEach(async () => {
+  for (const stream of streams.splice(0)) {
+    stream.close();
+    await stream.closed.catch(() => undefined);
+  }
+});
+
+function worker() {
+  const frames = fixtureEvents("notebook-events.sse");
+  const document = frames.find((event) => event.event === "document")?.data
+    .document;
+  assert(document);
+  const listeners = new Set<(event: unknown) => void>();
+  const reply = { ok: true, status: 200, data: { document, cells: {} } };
+  const wire = { status: 200, body: reply.data };
+  const bridge = {
+    notebookId: "n1",
+    subscribe: vi.fn(
+      async (_cursor: string | null, emit: (event: unknown) => void) => {
+        listeners.add(emit);
+        frames.forEach(emit);
+        return () => {
+          listeners.delete(emit);
+        };
+      },
+    ),
+    call: vi.fn<WasmNotebookBridge["call"]>().mockResolvedValue(wire),
+  };
+  return { bridge, listeners, reply, wire };
+}
+
+async function follow(bridge: WasmNotebookBridge) {
+  const stream = new NotebookEvents(
+    createWasmNotebookBackend(bridge).transport,
+  );
+  streams.push(stream);
+  const ready = new Deferred<void>();
+  stream.subscribe((event) => {
+    if (event.event === "ready") {
+      ready.resolve();
+    }
+  }, ready.reject);
+  await ready.promise;
+  return stream;
+}
+
+describe("WASM notebook backend", () => {
+  it("follows independent readers and releases each subscription on close", async () => {
+    const { bridge, listeners } = worker();
+    const first = await follow(bridge);
+    const second = await follow(bridge);
+    expect(listeners.size).toBe(2);
+    expect(second.status).toBe("connected");
+    first.close();
+    await first.closed;
+    expect(listeners.size).toBe(1);
+    second.close();
+    await second.closed;
+    expect(listeners.size).toBe(0);
+  });
+
+  it("closes during startup and releases the late subscription", async () => {
+    const { bridge } = worker();
+    const pending = new Deferred<() => void>();
+    bridge.subscribe.mockReturnValue(pending.promise);
+    const controller = new AbortController();
+    const done = createWasmNotebookBackend(bridge).transport(
+      { cursor: "previous-4", signal: controller.signal },
+      { open: vi.fn(), message: vi.fn() },
+    );
+    expect(bridge.subscribe).toHaveBeenCalledWith(
+      "previous-4",
+      expect.any(Function),
+    );
+    controller.abort();
+    await done;
+    const close = vi.fn();
+    pending.resolve(close);
+    await vi.waitFor(() => expect(close).toHaveBeenCalledOnce());
+  });
+
+  it("stops on worker initialization failure", async () => {
+    const { bridge } = worker();
+    bridge.subscribe.mockRejectedValue(new Error("Initialization failed"));
+    await expect(
+      createWasmNotebookBackend(bridge).transport(
+        { cursor: null, signal: new AbortController().signal },
+        { open: vi.fn(), message: vi.fn() },
+      ),
+    ).rejects.toBeInstanceOf(StreamStopped);
+  });
+
+  it("refuses commands for any notebook but the worker's", () => {
+    const { bridge } = worker();
+    expect(() => createWasmNotebookBackend(bridge).notebook("n2")).toThrow(
+      "serves notebook n1, not n2",
+    );
+    expect(bridge.call).not.toHaveBeenCalled();
+  });
+
+  it("sends edits as calls with the retry key and validates replies", async () => {
+    const { bridge, reply } = worker();
+    const backend = createWasmNotebookBackend(bridge);
+    expect(
+      await backend.notebook("n1").edit([], { idempotencyKey: "edit-7" }),
+    ).toEqual(reply);
+    expect(bridge.call).toHaveBeenCalledWith({
+      operation: "notebook.edit",
+      params: { notebook_id: "n1" },
+      headers: { "idempotency-key": "edit-7" },
+      body: { changes: [] },
+    });
+    const problem = {
+      type: "tag:marimo.io,2026:host/forbidden",
+      title: "Forbidden",
+      status: 403,
+    };
+    bridge.call.mockResolvedValue({ status: 403, body: problem });
+    expect(await backend.notebook("n1").edit([])).toEqual({
+      ok: false,
+      status: 403,
+      problem,
+    });
+    bridge.call.mockResolvedValue({ status: 200, body: {} });
+    await expect(backend.notebook("n1").edit([])).rejects.toThrow(
+      /Invalid input/,
+    );
+  });
+
+  it("cancels waiting for an accepted edit and permits retry with the same key", async () => {
+    const { bridge, reply, wire } = worker();
+    const pending = new Deferred<unknown>();
+    bridge.call.mockReturnValue(pending.promise);
+    const backend = createWasmNotebookBackend(bridge);
+    const controller = new AbortController();
+    const editing = backend.notebook("n1").edit([], {
+      signal: controller.signal,
+      idempotencyKey: "edit-8",
+    });
+    await vi.waitFor(() => expect(bridge.call).toHaveBeenCalledOnce());
+    controller.abort();
+    await expect(editing).rejects.toBe(controller.signal.reason);
+    pending.resolve(wire);
+    expect(
+      await backend.notebook("n1").edit([], { idempotencyKey: "edit-8" }),
+    ).toEqual(reply);
+    expect(bridge.call.mock.calls[0]).toEqual(bridge.call.mock.calls[1]);
+    await expect(
+      backend.notebook("n1").edit([], { signal: controller.signal }),
+    ).rejects.toBe(controller.signal.reason);
+    expect(bridge.call).toHaveBeenCalledTimes(2);
+  });
+});

package/src/core/host/backend.ts

@@ -0,0 +1,19 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A notebook's events and commands, from a host or a stand-in for one. */
+
+import type { NotebookClient, NotebookTransport } from "@marimo-team/client";
+import type { NotebookId } from "@marimo-team/client/protocol";
+
+/** Events and commands for a notebook, independent of React or a replica. */
+export interface NotebookBackend {
+  readonly transport: NotebookTransport;
+  /**
+   * Commands for a notebook the transport reported.
+   *
+   * Method presence is not permission. The notebook advertises its operations.
+   */
+  notebook(id: NotebookId): Pick<NotebookClient, "edit">;
+  /** Whether the host is reachable before an explicit reconnect. Absent means always. */
+  checkHealth?(): Promise<boolean>;
+}

package/src/core/host/connection.ts

@@ -37,45 +37,6 @@ export type ConnectionStatus =
 
 /** The recovery policy for a closed connection. */
 export interface CloseDecision {
-  kind: "terminal" | "gave-up" | "retry";
+  kind: "terminal" | "retry";
   status: ConnectionStatus;
 }
-
-export type ConnectionEvent = "open" | "close" | "message" | "error";
-
-export interface IConnectionTransportMap {
-  open: WebSocketEventMap["open"];
-  message: WebSocketEventMap["message"];
-  close: WebSocketEventMap["close"];
-  error: WebSocketEventMap["error"];
-}
-
-export type ConnectionTransportCallback<T extends ConnectionEvent> = (
-  data: IConnectionTransportMap[T],
-) => void;
-
-/** Transport readiness, using the WebSocket readyState values. */
-export const TransportReadyState = {
-  CONNECTING: 0,
-  OPEN: 1,
-  CLOSING: 2,
-  CLOSED: 3,
-} as const;
-
-export type TransportReadyState =
-  (typeof TransportReadyState)[keyof typeof TransportReadyState];
-
-/** A reconnectable source of notebook messages. */
-export interface NotebookTransport {
-  reconnect(code?: number | undefined, reason?: string | undefined): void;
-  close(): void;
-  addEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void;
-  removeEventListener<T extends ConnectionEvent>(
-    event: T,
-    callback: ConnectionTransportCallback<T>,
-  ): void;
-  readyState: TransportReadyState;
-}

package/src/core/host/index.ts

@@ -0,0 +1,14 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Notebook backends and the replica that follows them. */
+
+export type { NotebookBackend } from "./backend";
+export {
+  type CloseDecision,
+  ConnectionClosedReason,
+  ConnectionState,
+  type ConnectionStatus,
+} from "./connection";
+export { createNotebookReplica, type NotebookReplica } from "./replica";
+export type { Document } from "./resources/document";
+export { createWasmNotebookBackend } from "./wasm";

package/src/core/host/replica.ts

@@ -0,0 +1,286 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What a backend's stream has confirmed about a notebook, kept in atoms. */
+
+import { Logger } from "@/utils/Logger";
+import {
+  createNotebookReducer,
+  type NotebookEvent,
+  NotebookEvents,
+  StreamStopped,
+} from "@marimo-team/client";
+import type { Attachment, NotebookOutput } from "@marimo-team/client/protocol";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import { z } from "zod";
+import type { NotebookBackend } from "./backend";
+import {
+  ConnectionClosedReason,
+  ConnectionState,
+  type ConnectionStatus,
+} from "./connection";
+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 notebook as its backend's stream reports it, kept current while connected. */
+export interface NotebookReplica {
+  /** 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>;
+  /**
+   * The backend this replica follows, or null while not connected.
+   *
+   * Its commands reach the host that confirmed `documentAtom`. A temporary
+   * disconnection keeps it. Connecting again or disconnecting replaces it.
+   */
+  readonly backendAtom: Atom<NotebookBackend | null>;
+  /** Live attachments, empty until the stream is ready or while disconnected. */
+  readonly attachmentsAtom: Atom<readonly Attachment[]>;
+  /**
+   * Follow a backend, replacing any current connection.
+   *
+   * The returned function ends this connection and clears the replica.
+   * It does nothing once another connection has replaced this one.
+   */
+  connect(
+    backend: NotebookBackend,
+    options: {
+      /** Handle a native marimo notification. Its payload has not been validated. */
+      onNotification: (notification: unknown) => void;
+    },
+  ): () => void;
+  /** Restore the current connection after its stream closed, if the host is healthy. */
+  reconnect(): Promise<void>;
+}
+
+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 replica without connecting it. */
+export function createNotebookReplica(
+  store: ReturnType<typeof createStore>,
+): NotebookReplica {
+  // `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,
+  );
+  const backendAtom = atom<NotebookBackend | null>(null);
+  /** The connection this replica follows. Connecting again ends it. */
+  let current: { end(): void; reconnect(): Promise<void> } | undefined;
+
+  function setStatus(status: ConnectionStatus) {
+    store.set(readyAtom, false);
+    store.set(connectionAtom, status);
+  }
+
+  /** Drop what the stream reported and show `status`. */
+  function clear(status: ConnectionStatus) {
+    store.set(streamAtom, reducer.initial);
+    setStatus(status);
+  }
+
+  return {
+    connectionAtom,
+    notebookAtom,
+    documentAtom,
+    backendAtom,
+    attachmentsAtom,
+    connect,
+    reconnect: async () => current?.reconnect(),
+  };
+
+  function connect(
+    connected: NotebookBackend,
+    { onNotification }: { onNotification: (notification: unknown) => void },
+  ) {
+    current?.end();
+    const lifetime = new AbortController();
+    const { signal } = lifetime;
+    const checkHealth = connected.checkHealth ?? (async () => true);
+    store.set(backendAtom, connected);
+    const { transport } = connected;
+    let events: NotebookEvents<(typeof TYPES)[number]> | undefined;
+    let checkingHealth = false;
+    // Whether the replica still holds the state the last cursor describes.
+    let resumable = 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 stopMalformed(source: NotebookEvents<string>) {
+      source.close();
+      // The cursor can already include an event the replica rejected.
+      resumable = false;
+      clear(malformed);
+    }
+
+    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;
+      }
+      try {
+        store.set(streamAtom, reducer.reduce(store.get(streamAtom), event));
+      } catch (error) {
+        // 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 === "ready" || event.event === "reset") {
+        store.set(readyAtom, event.event === "ready");
+      }
+    }
+
+    const reconnect = async () => {
+      if (signal.aborted || checkingHealth || events?.status !== "closed") {
+        return;
+      }
+      checkingHealth = true;
+      setStatus({ state: ConnectionState.CONNECTING });
+      try {
+        const healthy = await checkHealth();
+        if (signal.aborted) {
+          return;
+        }
+        if (!healthy) {
+          setStatus(disconnected);
+          return;
+        }
+        follow(resumable ? (events.cursor ?? undefined) : undefined);
+      } catch (error) {
+        if (!signal.aborted) {
+          Logger.error("Notebook health check failed", error);
+          setStatus(disconnected);
+        }
+      } finally {
+        checkingHealth = false;
+      }
+    };
+
+    const connection = {
+      end: () => {
+        if (signal.aborted) {
+          return;
+        }
+        lifetime.abort();
+        if (current === connection) {
+          current = undefined;
+          store.set(backendAtom, null);
+        }
+        clear({ state: ConnectionState.NOT_STARTED });
+      },
+      reconnect,
+    };
+    current = connection;
+    setStatus({ state: ConnectionState.CONNECTING });
+    follow();
+    return connection.end;
+  }
+}