@marimo-team/frontend 0.25.3-dev2 → 0.25.3-dev20

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

@@ -0,0 +1,211 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import type { NotebookEvent } from "../events";
+import type { Resource } from "../resources/resource";
+import { createStream, type Stream, type StreamState } from "../stream";
+import { fixtureEvents } from "./fixtures";
+
+interface CellIds {
+  readonly ids: readonly string[];
+}
+
+/** A resource whose events are complete on their own. */
+const cells: Resource<CellIds> = {
+  initial: { ids: [] },
+  on: {
+    cell: (state, event) =>
+      state.ids.includes(event.data.cell.id)
+        ? state
+        : { ids: [...state.ids, event.data.cell.id] },
+    "cell.removed": (state, event) =>
+      state.ids.includes(event.data.id)
+        ? { ids: state.ids.filter((id) => id !== event.data.id) }
+        : state,
+  },
+};
+
+interface Revisions {
+  /** Cells changed since the last revision. */
+  readonly pending: readonly string[];
+  readonly published: readonly (readonly string[])[];
+}
+
+function publish(state: Revisions): Revisions {
+  return { pending: [], published: [...state.published, state.pending] };
+}
+
+/** A resource assembled from groups of events, like a document from its cells. */
+const revisions: Resource<Revisions> = {
+  initial: { pending: [], published: [] },
+  on: {
+    cell: (state, event) => ({
+      ...state,
+      pending: [...state.pending, event.data.cell.id],
+    }),
+    "cell.removed": (state, event) => ({
+      ...state,
+      pending: [...state.pending, event.data.id],
+    }),
+    // A live revision completes a group. A snapshot's revision waits for ready.
+    document: (state, _event, { phase }) =>
+      phase === "live" ? publish(state) : state,
+    ready: publish,
+  },
+};
+
+interface Resources {
+  cells: CellIds;
+  revisions: Revisions;
+}
+type State = StreamState<Resources>;
+
+function replay(
+  stream: Stream<Resources>,
+  events: readonly NotebookEvent[],
+  from: State = stream.initial,
+): State {
+  return events.reduce((state, event) => stream.reduce(state, event), from);
+}
+
+const snapshot = fixtureEvents("notebook-events.sse");
+const ready = snapshot.findIndex((event) => event.event === "ready");
+const cellEvent = snapshot.find((event) => event.event === "cell");
+if (cellEvent?.event !== "cell") {
+  throw new Error("Fixture is missing a cell event");
+}
+
+describe("notebook stream", () => {
+  it("reads the snapshot, goes live at ready, and rebuilds after reset", () => {
+    const stream = createStream({ cells, revisions });
+    const partial = replay(stream, snapshot.slice(0, ready));
+    expect(partial).toMatchObject({
+      phase: "snapshot",
+      notebookId: "n1",
+      cursor: null,
+    });
+    expect(partial.resources.cells.ids).toEqual(["c1", "c3", "k7f2"]);
+    expect(partial.resources.revisions).toEqual({
+      pending: ["c1", "c3", "k7f2"],
+      published: [],
+    });
+
+    const live = replay(stream, snapshot.slice(ready), partial);
+    expect(live).toMatchObject({ phase: "live", cursor: "97" });
+    expect(live.resources.cells).toBe(partial.resources.cells);
+    expect(live.resources.revisions).toEqual({
+      pending: [],
+      published: [["c1", "c3", "k7f2"]],
+    });
+
+    const resetEvents = fixtureEvents("events-reset.sse");
+    const afterReset = stream.reduce(live, resetEvents[0]);
+    expect(afterReset).toEqual({ ...stream.initial, notebookId: "n1" });
+    expect(replay(stream, resetEvents, live).phase).toBe("live");
+  });
+
+  it("keeps resource state by reference when an event does not apply", () => {
+    const stream = createStream({ cells, revisions });
+    const state = replay(stream, snapshot);
+    const runtime = snapshot.find((event) => event.event === "runtime");
+    const unhandled = stream.reduce(state, { ...runtime!, id: "100" });
+    expect(unhandled.resources).toBe(state.resources);
+    expect(unhandled.cursor).toBe("100");
+    expect(
+      stream.reduce(state, { event: "unknown", name: "x", id: "101" })
+        .resources,
+    ).toBe(state.resources);
+    expect(
+      stream.reduce(state, { event: "marimo", data: { marimo: {} } }),
+    ).toBe(state);
+  });
+
+  it("clears resources when its notebook is removed and ignores other notebooks", () => {
+    const stream = createStream({ cells, revisions });
+    const state = replay(stream, snapshot);
+    const time = "2026-10-01T09:30:00Z";
+    expect(
+      stream.reduce(state, {
+        event: "notebook.removed",
+        id: "100",
+        data: { id: "n2", time },
+      }).resources,
+    ).toBe(state.resources);
+    const removed = stream.reduce(state, {
+      event: "notebook.removed",
+      id: "100",
+      data: { id: "n1", time },
+    });
+    expect(removed).toMatchObject({ phase: "live", notebookId: "n1" });
+    expect(removed.resources).toBe(stream.initial.resources);
+  });
+
+  it("rejects objects from another notebook", () => {
+    const stream = createStream({ cells, revisions });
+    const state = replay(stream, snapshot.slice(0, ready));
+    const notebook = snapshot[0];
+    if (notebook.event !== "notebook") {
+      throw new Error("Fixture does not start with its notebook");
+    }
+    expect(() =>
+      stream.reduce(state, {
+        ...notebook,
+        data: { notebook: { ...notebook.data.notebook, id: "n2" } },
+      }),
+    ).toThrow("changed identity");
+    expect(() =>
+      stream.reduce(state, {
+        ...cellEvent,
+        data: { cell: { ...cellEvent.data.cell, notebook_id: "n2" } },
+      }),
+    ).toThrow("changed identity");
+  });
+
+  it("records each live event's cursor and resumes in place", () => {
+    const stream = createStream({ cells, revisions });
+    const partial = replay(stream, snapshot.slice(0, ready));
+    expect(stream.resume(partial)).toEqual({
+      ...stream.initial,
+      notebookId: "n1",
+    });
+
+    const live = replay(stream, snapshot.slice(ready), partial);
+    const [cell, , removed, published, saved] = fixtureEvents("edits.sse");
+    const editing = replay(stream, [cell], live);
+    expect(editing.cursor).toBe("110");
+    expect(editing.resources.revisions.pending).toEqual(["c2"]);
+    expect(stream.resume(editing)).toBe(editing);
+
+    const edited = replay(stream, [removed, published], editing);
+    expect(edited.cursor).toBe("113");
+    expect(edited.resources.revisions).toEqual({
+      pending: [],
+      published: [
+        ["c1", "c3", "k7f2"],
+        ["c2", "k7f2"],
+      ],
+    });
+    expect(stream.reduce(edited, saved).cursor).toBe("114");
+    expect(
+      stream.reduce(edited, { event: "ready", id: "116", data: {} }).cursor,
+    ).toBe("116");
+    expect(stream.reduce(edited, { event: "reset", data: {} }).cursor).toBe(
+      null,
+    );
+  });
+
+  it("requires a cursor on live events other than reset and notifications", () => {
+    const stream = createStream({ cells, revisions });
+    const state = replay(stream, snapshot);
+    expect(() => stream.reduce(state, cellEvent)).toThrow("cursor");
+    expect(
+      stream.reduce(state, { ...cellEvent, id: "100" }).resources.cells,
+    ).toBe(state.resources.cells);
+    expect(stream.reduce(state, { event: "reset", data: {} }).phase).toBe(
+      "snapshot",
+    );
+    expect(
+      stream.reduce(state, { event: "marimo", data: { marimo: null } }),
+    ).toBe(state);
+  });
+});

package/src/core/host/connection.ts

@@ -0,0 +1,81 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+export const ConnectionState = {
+  NOT_STARTED: "NOT_STARTED",
+  CONNECTING: "CONNECTING",
+  OPEN: "OPEN",
+  CLOSING: "CLOSING",
+  CLOSED: "CLOSED",
+} as const;
+
+export type ConnectionState =
+  (typeof ConnectionState)[keyof typeof ConnectionState];
+
+export const ConnectionClosedReason = {
+  KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
+  KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR",
+} as const;
+
+export type ConnectionClosedReason =
+  (typeof ConnectionClosedReason)[keyof typeof ConnectionClosedReason];
+
+export type ConnectionStatus =
+  | {
+      state: typeof ConnectionState.CLOSED;
+      code: ConnectionClosedReason;
+      reason: string;
+    }
+  | {
+      state: typeof ConnectionState.CONNECTING;
+    }
+  | {
+      state:
+        | typeof ConnectionState.OPEN
+        | typeof ConnectionState.CLOSING
+        | typeof ConnectionState.NOT_STARTED;
+    };
+
+/** The recovery policy for a closed connection. */
+export interface CloseDecision {
+  kind: "terminal" | "gave-up" | "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/events.ts

@@ -0,0 +1,60 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { z } from "zod";
+import * as Protocol from "./protocol";
+
+/**
+ * Notebook stream events, including native marimo notifications.
+ *
+ * Known events follow the generated protocol. Unknown events retain their name
+ * and cursor so clients can skip them without losing their place in the stream.
+ */
+
+/** Event kinds requested by the client; removals and controls are implicit. */
+export const NOTEBOOK_EVENT_TYPES = [
+  "notebook",
+  "document",
+  "cell",
+  "marimo",
+] as const satisfies readonly (Protocol.ObjectType | "marimo")[];
+
+const envelope = z.object({
+  event: z.string().min(1),
+  data: z.unknown(),
+  id: z.string().min(1).optional(),
+});
+
+const knownNames = new Set<string>(
+  Protocol.Event.options.flatMap((schema) => schema.shape.event.options),
+);
+
+/** A decoded event, or the name and cursor of an unrecognized event. */
+export type NotebookEvent =
+  | Protocol.EventOutput
+  | {
+      readonly event: "unknown";
+      readonly name: string;
+      readonly id?: string;
+    };
+
+const sessionEnvelope = z.object({ data: Protocol.MarimoNotification });
+
+/**
+ * Decode a host event or a legacy `{ data: notification }` session message.
+ *
+ * Throws if the envelope or a known event is malformed. Native notification
+ * payloads are passed through for the session handler to interpret.
+ */
+export function parseNotebookEvent(value: unknown): NotebookEvent {
+  if (typeof value === "object" && value !== null && !("event" in value)) {
+    return {
+      event: "marimo",
+      data: { marimo: sessionEnvelope.parse(value).data },
+    };
+  }
+  const parsed = envelope.parse(value);
+  if (knownNames.has(parsed.event)) {
+    return Protocol.Event.parse(parsed);
+  }
+  return { event: "unknown", name: parsed.event, id: parsed.id };
+}

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

@@ -0,0 +1,272 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { Logger } from "@/utils/Logger";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import {
+  type ConnectionStatus,
+  ConnectionState,
+  ConnectionClosedReason,
+  type CloseDecision,
+  type NotebookTransport,
+  TransportReadyState,
+} from "./connection";
+import { parseNotebookEvent } from "./events";
+import { type Document, document } from "./resources/document";
+import { createStream } from "./stream";
+
+/** A connection to a notebook on a marimo host. */
+export interface NotebookClient {
+  /** Client connection status, independent of application session readiness. */
+  readonly connectionAtom: Atom<ConnectionStatus>;
+  /**
+   * 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>;
+  /** Attempt to restore a closed connection after checking host health. */
+  reconnect(): Promise<void>;
+}
+
+/** The transport and application callbacks needed to connect 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. */
+  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 connectionAtom = atom<ConnectionStatus>({
+    state: ConnectionState.NOT_STARTED,
+  });
+  // One atom holds the whole stream, which resource atoms derive from.
+  const stream = createStream({ document });
+  const streamAtom = atom(stream.initial);
+  const documentAtom = atom(
+    (get) => get(streamAtom).resources.document.current,
+  );
+  let reconnect: (() => Promise<void>) | undefined;
+  const client: NotebookClient = {
+    connectionAtom,
+    documentAtom,
+    reconnect: async () => reconnect?.(),
+  };
+
+  function connect({
+    createTransport,
+    url,
+    messageFormat,
+    signal,
+    waitUntilReady,
+    checkHealth,
+    onNotification,
+    onMessageError,
+    decodeJson,
+    classifyClose,
+  }: NotebookConnection) {
+    if (signal.aborted) {
+      return;
+    }
+    let shouldRetry = true;
+    let terminal = false;
+    let checkingHealth = false;
+    let streamKey: string | undefined;
+    const transport = createTransport(() => {
+      if (messageFormat === "session") {
+        return url();
+      }
+      const next = new URL(url());
+      next.searchParams.delete("cursor");
+      if (signal.aborted) {
+        return next.toString();
+      }
+      // A cursor is useful with the state that produced it, which a change
+      // of session or notebook invalidates. Authentication may change freely.
+      const identity = new URL(next);
+      identity.searchParams.delete("access_token");
+      identity.searchParams.sort();
+      const key = identity.toString();
+      const state =
+        key === streamKey
+          ? stream.resume(store.get(streamAtom))
+          : stream.initial;
+      streamKey = key;
+      store.set(streamAtom, state);
+      if (state.cursor !== null) {
+        next.searchParams.set("cursor", state.cursor);
+      }
+      return next.toString();
+    });
+    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);
+      }
+    }
+
+    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;
+      }
+    }
+
+    function onError(event: Event) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      Logger.warn("Notebook connection error", event);
+      setStatus(disconnected);
+      retry();
+    }
+
+    function onMessage(message: MessageEvent<string>) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      let event;
+      try {
+        event = parseNotebookEvent(decodeJson(message.data));
+        store.set(streamAtom, stream.reduce(store.get(streamAtom), event));
+      } catch (error) {
+        Logger.error("Failed to decode notebook message", error);
+        if (messageFormat === "session") {
+          onMessageError(error);
+          return;
+        }
+        store.set(streamAtom, stream.initial);
+        terminal = true;
+        setStatus({
+          ...disconnected,
+          reason: "Failed to decode notebook message",
+        });
+        transport.close();
+        return;
+      }
+      if (event.event === "marimo") {
+        onNotification(event.data.marimo);
+      }
+    }
+
+    reconnect = async () => {
+      if (
+        signal.aborted ||
+        checkingHealth ||
+        transport.readyState === TransportReadyState.OPEN ||
+        transport.readyState === TransportReadyState.CONNECTING
+      ) {
+        return;
+      }
+      checkingHealth = true;
+      setStatus({ state: ConnectionState.CONNECTING });
+      try {
+        const healthy = await checkHealth();
+        if (signal.aborted) {
+          return;
+        }
+        if (!healthy) {
+          shouldRetry = false;
+          setStatus(disconnected);
+          return;
+        }
+        terminal = false;
+        shouldRetry = true;
+        transport.reconnect();
+      } catch (error) {
+        if (!signal.aborted) {
+          Logger.error("Notebook health check failed", error);
+          shouldRetry = false;
+          setStatus(disconnected);
+        }
+      } finally {
+        checkingHealth = false;
+      }
+    };
+
+    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, stream.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();
+          }
+        });
+    }
+  }
+
+  return { client, connect };
+}