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

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

@@ -8,13 +8,30 @@ import { Deferred } from "@/utils/Deferred";
 import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { classifyCloseEvent } from "../../websocket/connection-state";
 import { createNotebookClient } from "../notebook-client";
+import { fixtureMessages } from "./fixtures";
 
-function setup() {
+function notebookMessage(id: string) {
+  const notebook = {
+    id,
+    project_id: "p1",
+    path: null,
+    name: id,
+    operations: [],
+    updated_at: null,
+  };
+  return JSON.stringify({ event: "notebook", data: { notebook } });
+}
+
+function setup(url: () => string = () => "ws://localhost/ws") {
   const events = new EventTarget();
   const healthy = new Deferred<void>();
+  let provider = url;
+  const urls: URL[] = [];
   const transport = {
     readyState: WebSocket.CLOSED,
-    reconnect: vi.fn(),
+    reconnect: vi.fn(() => {
+      urls.push(new URL(provider()));
+    }),
     close: vi.fn(() => events.dispatchEvent(new CloseEvent("close"))),
     send: vi.fn(),
     addEventListener: events.addEventListener.bind(events),
@@ -27,8 +44,11 @@ function setup() {
   const client = controller.client;
   const connect = (signal: AbortSignal) =>
     controller.connect({
-      createTransport: () => transport,
-      url: () => "ws://localhost/ws",
+      createTransport: (next) => {
+        provider = next;
+        return transport;
+      },
+      url,
       messageFormat: "host",
       signal,
       waitUntilReady: () => healthy.promise,
@@ -47,6 +67,11 @@ function setup() {
     onNotification,
     store,
     checkHealth,
+    urls,
+    receive: (message: unknown) =>
+      events.dispatchEvent(
+        new MessageEvent("message", { data: JSON.stringify(message) }),
+      ),
   };
 }
 
@@ -79,6 +104,93 @@ describe("NotebookClient lifetime", () => {
     },
   );
 
+  it("resumes from the last cursor until reset or a new session", async () => {
+    const session = { id: "first", token: "a" };
+    const { connect, events, healthy, urls, receive } = setup(
+      () =>
+        `ws://localhost/ws?session_id=${session.id}&access_token=${session.token}&cursor=999`,
+    );
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    healthy.resolve();
+    await healthy.promise;
+    expect(urls[0].searchParams.get("cursor")).toBeNull();
+
+    const snapshot = fixtureMessages("notebook-events.sse");
+    events.dispatchEvent(new Event("open"));
+    for (const message of snapshot) {
+      receive(message);
+    }
+    session.token = "b";
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(urls[1].searchParams.get("cursor")).toBe("97");
+    expect(urls[1].searchParams.get("access_token")).toBe("b");
+
+    events.dispatchEvent(new Event("open"));
+    receive({ event: "reset", data: {} });
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(urls[2].searchParams.get("cursor")).toBeNull();
+
+    events.dispatchEvent(new Event("open"));
+    for (const message of snapshot) {
+      receive(message);
+    }
+    session.id = "second";
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(urls[3].searchParams.get("cursor")).toBeNull();
+    expect(urls[3].searchParams.get("session_id")).toBe("second");
+    lifetime.abort();
+  });
+
+  it("exposes the confirmed document until the lifetime ends", async () => {
+    const { connect, events, client, store, healthy, receive } = setup();
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    healthy.resolve();
+    await healthy.promise;
+    events.dispatchEvent(new Event("open"));
+    const snapshot = fixtureMessages("notebook-events.sse");
+    for (const message of snapshot.slice(0, -3)) {
+      receive(message);
+    }
+    expect(store.get(client.documentAtom)).toBeNull();
+    for (const message of snapshot.slice(-3)) {
+      receive(message);
+    }
+    expect(store.get(client.documentAtom)?.revision.order).toEqual([
+      "c1",
+      "c3",
+      "k7f2",
+    ]);
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(store.get(client.documentAtom)?.revision.order).toEqual([
+      "c1",
+      "c3",
+      "k7f2",
+    ]);
+    lifetime.abort();
+    expect(store.get(client.documentAtom)).toBeNull();
+  });
+
+  it("stops consuming a stream whose notebook changes", () => {
+    const { connect, events, client, store, transport } = setup();
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    events.dispatchEvent(new Event("open"));
+    for (const id of ["n1", "n2"]) {
+      events.dispatchEvent(
+        new MessageEvent("message", { data: notebookMessage(id) }),
+      );
+    }
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.CLOSED,
+      code: "KERNEL_DISCONNECTED",
+      reason: "Failed to decode notebook message",
+    });
+    expect(transport.close).toHaveBeenCalledOnce();
+    lifetime.abort();
+  });
+
   it.each(["MARIMO_UNAUTHORIZED", "MARIMO_TRANSPORT_EXHAUSTED"])(
     "stops automatic recovery after %s but permits an explicit reconnect",
     async (reason) => {

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/notebook-client.ts

@@ -11,11 +11,21 @@ import {
   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>;
 }
@@ -49,9 +59,16 @@ 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?.(),
   };
 
@@ -73,7 +90,33 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
     let shouldRetry = true;
     let terminal = false;
     let checkingHealth = false;
-    const transport = createTransport(url);
+    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 = {
@@ -130,12 +173,14 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
       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,
@@ -193,6 +238,7 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
         transport.removeEventListener("message", onMessage);
         transport.close();
         reconnect = undefined;
+        store.set(streamAtom, stream.initial);
         setStatus({ state: ConnectionState.NOT_STARTED });
       },
       { once: true },

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

@@ -0,0 +1,101 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "../protocol";
+import type { Resource } from "./resource";
+
+/**
+ * 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.
+ */
+export 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: Resource<DocumentState> = {
+  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/resources/resource.ts

@@ -0,0 +1,59 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "../protocol";
+
+/** An event name from the protocol. */
+export type EventName = Protocol.EventOutput["event"];
+
+/** The decoded event with the given name. */
+export type EventOf<Name extends EventName> = Extract<
+  Protocol.EventOutput,
+  { event: Name }
+>;
+
+/**
+ * Events the stream handles itself.
+ *
+ * They identify the notebook, discard state, or carry a native
+ * notification. Resources do not receive them.
+ */
+export type StreamEventName =
+  | "reset"
+  | "notebook"
+  | "notebook.removed"
+  | "marimo";
+
+/** Events a resource may handle. `ready` marks the end of the snapshot. */
+export type ResourceEventName = Exclude<EventName, StreamEventName>;
+
+/** `snapshot` until `ready` arrives and `live` after it. */
+export type StreamPhase = "snapshot" | "live";
+
+/** What a handler may know about the stream beyond the event. */
+export interface HandlerContext {
+  readonly phase: StreamPhase;
+}
+
+/**
+ * A pure handler for each event the resource reads.
+ *
+ * A handler returns `state` when the event does not apply to it.
+ */
+export type EventHandlers<State> = {
+  readonly [Name in ResourceEventName]?: (
+    state: State,
+    event: EventOf<Name>,
+    context: HandlerContext,
+  ) => State;
+};
+
+/**
+ * One kind of object assembled from its events.
+ *
+ * The stream returns a resource to `initial` on `reset` and when the
+ * notebook is removed.
+ */
+export interface Resource<State> {
+  readonly initial: State;
+  readonly on: EventHandlers<State>;
+}