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

package/src/core/edit-app.tsx

@@ -41,7 +41,7 @@ import { viewStateAtom } from "./mode";
 import { useRequestClient } from "./network/requests";
 import { useFilename } from "./saving/filename";
 import { lastSavedNotebookAtom } from "./saving/state";
-import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
+import { useNotebookSession } from "./websocket/useNotebookSession";
 
 interface AppProps {
   /**
@@ -83,7 +83,7 @@ export const EditApp: React.FC<AppProps> = ({
     };
   }, []);
 
-  const { connection, reconnect } = useMarimoKernelConnection({
+  const { connection, reconnect } = useNotebookSession({
     autoInstantiate: userConfig.runtime.auto_instantiate,
     setCells: (cells, layout) => {
       setCells(cells);

package/src/core/host/AGENTS.md

@@ -4,3 +4,8 @@ Production code may import host modules, `jotai/vanilla`, `zod`, and
 `@/utils/Logger`. Enforced by `marimo/host-imports`. Tests are exempt.
 
 Define interfaces here and inject application behavior through adapters.
+
+`stream.ts` reduces stream events into one state. Each kind of object is a
+`Resource` in `resources/` with pure handlers keyed by event name. The stream
+handles phase, notebook identity, and control events itself. The client holds
+one state atom and exposes derived atoms.

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

@@ -0,0 +1,124 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import type { NotebookEvent } from "../events";
+import { document } from "../resources/document";
+import { createStream } from "../stream";
+import { fixtureEvents } from "./fixtures";
+
+const stream = createStream({ document });
+const snapshot = fixtureEvents("notebook-events.sse");
+const edits = fixtureEvents("edits.sse");
+
+function replay(events: readonly NotebookEvent[], from = stream.initial) {
+  return events.reduce((state, event) => stream.reduce(state, event), from);
+}
+
+const confirmed = () => replay(snapshot);
+const current = (state: ReturnType<typeof replay>) =>
+  state.resources.document.current;
+
+describe("document resource", () => {
+  it("publishes the snapshot at ready, then each completed edit and save", () => {
+    const ready = snapshot.findIndex((event) => event.event === "ready");
+    const opening = replay(snapshot.slice(0, ready));
+    expect(current(opening)).toBeNull();
+    const first = current(replay(snapshot.slice(ready), opening));
+    expect(first?.revision.order).toEqual(["c1", "c3", "k7f2"]);
+    expect([...first!.cells.keys()]).toEqual(first?.revision.order);
+
+    const [edit, save] = edits.filter((event) => event.event === "document");
+    const pending = replay(
+      edits.filter((event) => event.event !== "document"),
+      confirmed(),
+    );
+    expect(current(pending)).toEqual(first);
+    const changed = stream.reduce(pending, edit);
+    expect(current(changed)?.revision).toEqual(edit.data.document);
+    expect([...current(changed)!.cells.keys()].toSorted()).toEqual(
+      edit.data.document.order.toSorted(),
+    );
+    const saved = current(stream.reduce(changed, save));
+    expect(saved?.revision).toEqual(save.data.document);
+    expect(saved?.cells.get("c3")?.code).toBe("x = 22");
+    expect(saved?.cells.has("k7f2")).toBe(false);
+    expect(saved?.revision.saved_etag).toBe(saved?.revision.etag);
+    expect(first?.cells.get("c3")?.code).toBe("x = 21");
+    expect(first?.cells.has("k7f2")).toBe(true);
+  });
+
+  it("applies order-only changes and distinguishes an empty notebook from no document", () => {
+    const state = confirmed();
+    const revision = current(state)!.revision;
+    const reordered = {
+      ...revision,
+      order: revision.order.toReversed(),
+      version: revision.version + 1,
+    };
+    const changed = stream.reduce(state, {
+      event: "document",
+      data: { document: reordered },
+      id: "200",
+    });
+    expect(current(changed)).toEqual({
+      revision: reordered,
+      cells: current(state)?.cells,
+    });
+
+    const empty = { ...revision, order: [] };
+    const opening = stream.reduce(stream.initial, {
+      event: "document",
+      data: { document: empty },
+    });
+    expect(current(opening)).toBeNull();
+    expect(
+      current(stream.reduce(opening, { event: "ready", data: {}, id: "0" })),
+    ).toEqual({ revision: empty, cells: new Map() });
+    expect(
+      current(
+        stream.reduce(stream.initial, { event: "ready", data: {}, id: "0" }),
+      ),
+    ).toBeNull();
+  });
+
+  it("keeps the document while a live update is pending and after a replay completes it", () => {
+    const state = confirmed();
+    const [cell, , removed, published] = edits;
+    if (published.event !== "document") {
+      throw new Error("Fixture does not publish a document");
+    }
+    const partial = replay([cell, removed], state);
+    expect(current(partial)).toBe(current(state));
+    expect(stream.resume(partial)).toBe(partial);
+    const completed = replay(
+      [published, { event: "ready", data: {}, id: "113" }],
+      partial,
+    );
+    expect(current(completed)?.revision).toEqual(published.data.document);
+  });
+
+  it("removes a notebook without waiting for individual cell removals", () => {
+    const removed = stream.reduce(confirmed(), {
+      event: "notebook.removed",
+      data: { id: "n1", time: "2026-10-01T09:30:00Z" },
+      id: "200",
+    });
+    expect(current(removed)).toBeNull();
+  });
+
+  it("rejects records that cannot form a complete revision", () => {
+    const state = confirmed();
+    const revision = current(state)!.revision;
+    expect(() =>
+      stream.reduce(state, {
+        event: "document",
+        data: { document: { ...revision, order: ["missing"] } },
+        id: "200",
+      }),
+    ).toThrow("Document records do not match");
+    const partial = stream.reduce(state, edits[0]);
+    expect(() =>
+      stream.reduce(partial, { event: "ready", data: {}, id: "200" }),
+    ).toThrow("incomplete document update");
+  });
+});

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

@@ -0,0 +1,46 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { z } from "zod";
+import { parseNotebookEvent } from "../events";
+
+describe("notebook messages", () => {
+  it("normalizes session and host envelopes for native marimo messages", () => {
+    const notification = {
+      op: "variables",
+      variables: [{ name: "x", value: Number.NaN }],
+    };
+    const host = { event: "marimo", data: { marimo: notification } };
+    expect(parseNotebookEvent({ data: notification })).toEqual(host);
+    expect(parseNotebookEvent({ ...host, id: "42" })).toEqual({
+      ...host,
+      id: "42",
+    });
+  });
+
+  it("retains unknown event cursors without interpreting their payloads", () => {
+    expect(
+      parseNotebookEvent({ event: "future", id: "42", data: null }),
+    ).toEqual({ event: "unknown", name: "future", id: "42" });
+  });
+
+  it("decodes known events outside the client's subscription filters", () => {
+    const event = {
+      event: "run.removed",
+      id: "42",
+      data: { id: "r1", time: "2026-10-19T12:00:00Z" },
+    };
+    expect(parseNotebookEvent(event)).toEqual(event);
+  });
+
+  it.each([
+    { event: "cell", id: "42", data: { cell: { id: "c1" } } },
+    { event: "run.removed", id: "42", data: {} },
+    { event: "ready", data: {} },
+    { event: "marimo", data: {} },
+    {},
+    null,
+  ])("rejects malformed known messages: %j", (message) => {
+    expect(() => parseNotebookEvent(message)).toThrow(z.ZodError);
+  });
+});

package/src/core/host/__tests__/fixtures.ts

@@ -0,0 +1,39 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import NodeFs from "node:fs";
+import NodeUrl from "node:url";
+import { parseNotebookEvent, type NotebookEvent } from "../events";
+
+/** Read the protocol's portable examples as transport messages. */
+export function fixtureMessages(name: string) {
+  const path = new NodeUrl.URL(
+    `../../../../../protocols/host/fixtures/${name}`,
+    import.meta.url,
+  );
+  return NodeFs.readFileSync(path, "utf8")
+    .split(/\r?\n\r?\n/)
+    .flatMap((frame) => {
+      const event = /^event: (.+)$/m.exec(frame)?.[1];
+      if (!event) {
+        return [];
+      }
+      const data = frame
+        .split(/\r?\n/)
+        .filter((line) => line.startsWith("data:"))
+        .map((line) => line.slice(5).trimStart())
+        .join("\n");
+      const payload: unknown = JSON.parse(data);
+      return [
+        {
+          event,
+          id: /^id: (.+)$/m.exec(frame)?.[1],
+          data: payload,
+        },
+      ];
+    });
+}
+
+/** Read the protocol's portable examples as decoded notebook events. */
+export function fixtureEvents(name: string): NotebookEvent[] {
+  return fixtureMessages(name).map(parseNotebookEvent);
+}

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

@@ -0,0 +1,303 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { createStore } from "jotai";
+import { ConnectionState } from "../../websocket/types";
+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 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(() => {
+      urls.push(new URL(provider()));
+    }),
+    close: vi.fn(() => events.dispatchEvent(new CloseEvent("close"))),
+    send: vi.fn(),
+    addEventListener: events.addEventListener.bind(events),
+    removeEventListener: events.removeEventListener.bind(events),
+  };
+  const onNotification = vi.fn();
+  const checkHealth = vi.fn(async () => true);
+  const store = createStore();
+  const controller = createNotebookClient(store);
+  const client = controller.client;
+  const connect = (signal: AbortSignal) =>
+    controller.connect({
+      createTransport: (next) => {
+        provider = next;
+        return transport;
+      },
+      url,
+      messageFormat: "host",
+      signal,
+      waitUntilReady: () => healthy.promise,
+      checkHealth,
+      onNotification: (notification) => onNotification({ data: notification }),
+      onMessageError: vi.fn(),
+      decodeJson: jsonParseWithSpecialChar,
+      classifyClose: classifyCloseEvent,
+    });
+  return {
+    events,
+    healthy,
+    transport,
+    client,
+    connect,
+    onNotification,
+    store,
+    checkHealth,
+    urls,
+    receive: (message: unknown) =>
+      events.dispatchEvent(
+        new MessageEvent("message", { data: JSON.stringify(message) }),
+      ),
+  };
+}
+
+describe("NotebookClient lifetime", () => {
+  it.each(["{broken", JSON.stringify({ event: "ready", data: {} })])(
+    "stops consuming a malformed host stream: %s",
+    (data) => {
+      const { connect, events, client, store, transport, onNotification } =
+        setup();
+      const lifetime = new AbortController();
+      connect(lifetime.signal);
+      events.dispatchEvent(new Event("open"));
+      events.dispatchEvent(new MessageEvent("message", { data }));
+      events.dispatchEvent(
+        new MessageEvent("message", {
+          data: JSON.stringify({
+            event: "marimo",
+            data: { marimo: { op: "reconnected" } },
+          }),
+        }),
+      );
+      expect(store.get(client.connectionAtom)).toEqual({
+        state: ConnectionState.CLOSED,
+        code: "KERNEL_DISCONNECTED",
+        reason: "Failed to decode notebook message",
+      });
+      expect(transport.close).toHaveBeenCalledOnce();
+      expect(onNotification).not.toHaveBeenCalled();
+      lifetime.abort();
+    },
+  );
+
+  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) => {
+      const {
+        connect,
+        client,
+        store,
+        events,
+        transport,
+        healthy,
+        checkHealth,
+      } = setup();
+      const lifetime = new AbortController();
+      connect(lifetime.signal);
+      healthy.resolve();
+      await healthy.promise;
+      events.dispatchEvent(new CloseEvent("close", { reason }));
+      events.dispatchEvent(new Event("error"));
+      expect(store.get(client.connectionAtom).state).toBe(
+        ConnectionState.CLOSED,
+      );
+      expect(transport.reconnect).toHaveBeenCalledOnce();
+      await client.reconnect();
+      expect(checkHealth).toHaveBeenCalledOnce();
+      expect(transport.reconnect).toHaveBeenCalledTimes(2);
+      lifetime.abort();
+    },
+  );
+
+  it("only resets the automatic retry budget after an open event", () => {
+    const { connect, events, transport } = setup();
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    events.dispatchEvent(new CloseEvent("close"));
+    events.dispatchEvent(new Event("error"));
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+    events.dispatchEvent(new Event("open"));
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(transport.reconnect).toHaveBeenCalledTimes(2);
+    lifetime.abort();
+  });
+
+  it("does not revive an aborted lifetime when a manual health check finishes", async () => {
+    const { connect, client, transport, checkHealth } = setup();
+    const health = new Deferred<boolean>();
+    checkHealth.mockReturnValue(health.promise);
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    const attempt = client.reconnect();
+    lifetime.abort();
+    health.resolve(true);
+    await attempt;
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("does not connect after abort, even when the health check finishes", async () => {
+    const { client, connect, healthy, transport } = setup();
+    const controller = new AbortController();
+    connect(controller.signal);
+    controller.abort();
+    healthy.resolve();
+    await healthy.promise;
+    client.reconnect();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(transport.close).toHaveBeenCalledOnce();
+  });
+
+  it("detaches listeners before closing so cleanup cannot trigger a session retry", () => {
+    const { connect, client, store, events, transport, onNotification } =
+      setup();
+    const controller = new AbortController();
+    connect(controller.signal);
+    controller.abort();
+    events.dispatchEvent(
+      new MessageEvent("message", {
+        data: JSON.stringify({ data: { op: "reconnected" } }),
+      }),
+    );
+    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.NOT_STARTED,
+    });
+    expect(onNotification).not.toHaveBeenCalled();
+  });
+
+  it("can reconnect after effect cleanup without reviving the earlier client attempt", async () => {
+    const { connect, healthy, transport, events, onNotification } = setup();
+    const first = new AbortController();
+    connect(first.signal);
+    first.abort();
+    const second = new AbortController();
+    connect(second.signal);
+    healthy.resolve();
+    await healthy.promise;
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+    events.dispatchEvent(
+      new MessageEvent("message", {
+        data: JSON.stringify({
+          event: "marimo",
+          data: { marimo: { op: "reconnected" } },
+        }),
+      }),
+    );
+    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
+      data: { op: "reconnected" },
+    });
+    second.abort();
+  });
+});