@marimo-team/frontend 0.25.3-dev79 → 0.25.3-dev81

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

@@ -1,27 +1,31 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { Logger } from "@/utils/Logger";
+import {
+  createNotebookReducer,
+  type NotebookEvent,
+  NotebookEvents,
+  type NotebookTransport,
+  StreamStopped,
+} from "@marimo-team/client";
 import type { Attachment, NotebookOutput } from "@marimo-team/client/protocol";
-import { createNotebookReducer } from "@marimo-team/client";
 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 { type NotebookEvent, parseNotebookEvent } from "./events";
 import { attachments } from "./resources/attachments";
 import { type Document, document } from "./resources/document";
 
-/** Events that may arrive without a cursor after `ready`. */
-const UNCURSORED = new Set<NotebookEvent["event"]>([
-  "reset",
-  "marimo",
-  "unknown",
-]);
+/** 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 {
@@ -43,28 +47,53 @@ export interface NotebookClient {
   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;
-  /** 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;
-  /** Choose the connection status and retry policy after the transport closes. */
-  classifyClose: (event: CloseEvent) => CloseDecision;
+}
+
+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. */
 export function createNotebookClient(store: ReturnType<typeof createStore>) {
+  // `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,
   });
@@ -73,16 +102,18 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
   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(readyAtom) && get(connectionAtom).state === ConnectionState.OPEN
-      ? get(streamAtom).notebook
-      : null,
+    get(presentAtom) ? get(streamAtom).notebook : null,
   );
   const documentAtom = atom(
     (get) => get(streamAtom).resources.document.current,
   );
   const attachmentsAtom = atom((get) =>
-    get(readyAtom) && get(connectionAtom).state === ConnectionState.OPEN
+    get(presentAtom)
       ? get(streamAtom).resources.attachments
       : attachments.initial,
   );
@@ -95,156 +126,97 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
     reconnect: async () => reconnect?.(),
   };
 
+  function setStatus(status: ConnectionStatus) {
+    store.set(readyAtom, false);
+    store.set(connectionAtom, status);
+  }
+
   function connect({
-    createTransport,
-    url,
+    transport,
     signal,
-    waitUntilReady,
     checkHealth,
     onNotification,
-    decodeJson,
-    classifyClose,
   }: NotebookConnection) {
     if (signal.aborted) {
       return;
     }
-    let shouldRetry = true;
-    let terminal = false;
+    let events: NotebookEvents<(typeof TYPES)[number]> | undefined;
     let checkingHealth = false;
-    let streamKey: string | undefined;
-    // The id of the last live event. Null before `ready` and after `reset`.
-    let cursor: string | null = null;
-    const transport = createTransport(() => {
-      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 current = store.get(streamAtom);
-      // A live stream keeps its state and resumes from the cursor. An
-      // incomplete snapshot starts over.
-      const state =
-        key !== streamKey || current.phase === "snapshot"
-          ? reducer.initial
-          : current;
-      if (state.phase === "snapshot") {
-        cursor = null;
-      }
-      streamKey = key;
-      store.set(readyAtom, false);
-      store.set(streamAtom, state);
-      if (cursor !== null) {
-        next.searchParams.set("cursor", cursor);
-      }
-      return next.toString();
-    });
-    const setStatus = (status: ConnectionStatus) => {
-      store.set(readyAtom, false);
-      store.set(connectionAtom, status);
-    };
-    const disconnected = {
-      state: ConnectionState.CLOSED,
-      code: ConnectionClosedReason.KERNEL_DISCONNECTED,
-      reason: "kernel not found",
-    } satisfies ConnectionStatus;
+    // Whether the replica still holds the state the last cursor describes.
+    let resumable = true;
 
-    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;
+    /** 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));
-        const state = store.get(streamAtom);
-        if (
-          state.phase === "live" &&
-          event.id === undefined &&
-          !UNCURSORED.has(event.event)
-        ) {
-          throw new Error("Live notebook event is missing its cursor");
-        }
-        if (event.event !== "marimo" && event.event !== "unknown") {
-          store.set(streamAtom, reducer.reduce(state, event));
-        }
-        if (event.event === "reset") {
-          cursor = null;
-        } else if (event.event === "ready" || state.phase === "live") {
-          cursor = event.id ?? cursor;
-        }
-        if (event.event === "ready" || event.event === "reset") {
-          store.set(readyAtom, event.event === "ready");
-        }
+        store.set(streamAtom, reducer.reduce(store.get(streamAtom), event));
       } catch (error) {
-        Logger.error("Failed to decode notebook message", error);
-        store.set(streamAtom, reducer.initial);
-        cursor = null;
-        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;
@@ -255,17 +227,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 {
@@ -276,11 +244,6 @@ 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 });
@@ -288,28 +251,8 @@ export function createNotebookClient(store: ReturnType<typeof createStore>) {
       { 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/websocket/__tests__/useNotebookClient.hook.test.tsx

@@ -5,52 +5,48 @@ import { act, renderHook } from "@testing-library/react";
 import { createStore, Provider } from "jotai";
 import type { PropsWithChildren } from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { NotebookTransport } from "@marimo-team/client";
 import { ConnectionState } from "../types";
-import { createConnectionTransport } from "../transports/create";
-import type { IConnectionTransport } from "../transports/transport";
+import { Deferred } from "@/utils/Deferred";
+import { createNotebookTransport } from "../transports/notebook";
 import { useNotebookClient } from "../useNotebookClient";
 import { fixtureMessages } from "../../host/__tests__/fixtures";
 import { agentAttachmentsAtom } from "../../attachments/state";
 import { notebookClientAtom } from "../state";
 
-vi.mock("../transports/create", () => ({ createConnectionTransport: vi.fn() }));
+vi.mock("../transports/notebook", () => ({ createNotebookTransport: vi.fn() }));
 
-function makeTransport(url?: () => string) {
-  const events = new EventTarget();
-  const transport: IConnectionTransport = {
-    readyState: WebSocket.CLOSED,
-    reconnect: vi.fn(() => {
-      url?.();
-      transport.readyState = WebSocket.OPEN;
-      events.dispatchEvent(new Event("open"));
-    }),
-    close: vi.fn(() => {
-      transport.readyState = WebSocket.CLOSED;
-      events.dispatchEvent(new CloseEvent("close"));
-    }),
-    send: vi.fn(),
-    addEventListener: events.addEventListener.bind(events),
-    removeEventListener: events.removeEventListener.bind(events),
-  };
+/** A transport whose connections the test drives. */
+function makeTransport() {
+  const opened: {
+    signal: AbortSignal;
+    open(): void;
+    receive(message: unknown): void;
+  }[] = [];
+  const transport: NotebookTransport = ({ signal }, handlers) =>
+    new Promise((resolve) => {
+      opened.push({ signal, open: handlers.open, receive: handlers.message });
+      signal.addEventListener("abort", () => resolve(), { once: true });
+    });
   return {
     transport,
-    receive: (message: unknown) =>
-      events.dispatchEvent(
-        new MessageEvent("message", { data: JSON.stringify(message) }),
-      ),
+    opened,
+    /** Open the latest connection and deliver a host-framed notification. */
     notify: () =>
-      events.dispatchEvent(
-        new MessageEvent("message", {
-          data: JSON.stringify({
-            event: "marimo",
-            data: { marimo: { op: "reconnected" } },
-          }),
-        }),
-      ),
+      opened.at(-1)?.receive({
+        event: "marimo",
+        data: { marimo: { op: "reconnected" } },
+      }),
   };
 }
 
-function setup(strict = false) {
+function setup(
+  strict = false,
+  readiness: {
+    waitUntilReady?: () => Promise<void>;
+    checkHealth?: () => Promise<boolean>;
+  } = {},
+) {
   const store = createStore();
   const options = {
     static: false,
@@ -60,6 +56,7 @@ function setup(strict = false) {
     waitUntilReady: async () => undefined,
     checkHealth: async () => true,
     onNotification: vi.fn(),
+    ...readiness,
   };
   const wrapper = ({ children }: PropsWithChildren) => (
     <Provider store={store}>{children}</Provider>
@@ -73,19 +70,97 @@ function setup(strict = false) {
 }
 
 beforeEach(() => {
-  vi.mocked(createConnectionTransport).mockReset();
+  vi.mocked(createNotebookTransport).mockReset();
 });
 
 describe("useNotebookClient", () => {
+  it.each([true, false])(
+    "rechecks health after failed startup (healthy: %s)",
+    async (healthy) => {
+      const fake = makeTransport();
+      vi.mocked(createNotebookTransport).mockImplementation(
+        ({ url }) =>
+          async (attempt, handlers) => {
+            if (typeof url === "function") {
+              await url();
+            }
+            if (!attempt.signal.aborted) {
+              return fake.transport(attempt, handlers);
+            }
+          },
+      );
+      const waitUntilReady = vi
+        .fn()
+        .mockRejectedValue(new Error("Startup gave up"));
+      const checkHealth = vi.fn().mockResolvedValue(healthy);
+      const { result, store, unmount } = setup(false, {
+        waitUntilReady,
+        checkHealth,
+      });
+      await act(async () => Promise.resolve());
+      expect(store.get(result.current.connectionAtom).state).toBe(
+        ConnectionState.CLOSED,
+      );
+
+      await act(async () => result.current.reconnect());
+      expect({
+        healthChecks: checkHealth.mock.calls.length,
+        initialWaits: waitUntilReady.mock.calls.length,
+        connections: fake.opened.length,
+        state: store.get(result.current.connectionAtom).state,
+      }).toEqual({
+        healthChecks: 1,
+        initialWaits: 1,
+        connections: healthy ? 1 : 0,
+        state: healthy ? ConnectionState.CONNECTING : ConnectionState.CLOSED,
+      });
+      act(() => fake.opened.at(-1)?.open());
+      expect(store.get(result.current.connectionAtom).state).toBe(
+        healthy ? ConnectionState.OPEN : ConnectionState.CLOSED,
+      );
+      unmount();
+    },
+  );
+
+  it("waits for initial readiness without reviving an unmounted connection", async () => {
+    const ready = new Deferred<void>();
+    const resolvedUrl = vi.fn();
+    vi.mocked(createNotebookTransport).mockImplementation(
+      ({ url }) =>
+        async () => {
+          if (typeof url === "function") {
+            await url();
+            resolvedUrl();
+          }
+        },
+    );
+    const { result, store, unmount } = setup(false, {
+      waitUntilReady: () => ready.promise,
+    });
+    await act(async () => Promise.resolve());
+    expect(resolvedUrl).not.toHaveBeenCalled();
+    expect(store.get(result.current.connectionAtom)).toEqual({
+      state: ConnectionState.CONNECTING,
+    });
+    unmount();
+    await act(async () => ready.resolve());
+    expect(store.get(result.current.connectionAtom)).toEqual({
+      state: ConnectionState.NOT_STARTED,
+    });
+    expect(store.get(notebookClientAtom)).toBeNull();
+  });
+
   it("exposes host presence to the UI without waiting for kernel-ready", async () => {
-    const transport = makeTransport();
-    vi.mocked(createConnectionTransport).mockReturnValue(transport.transport);
+    const fake = makeTransport();
+    vi.mocked(createNotebookTransport).mockReturnValue(fake.transport);
     const { result, store, unmount } = setup();
     await act(async () => Promise.resolve());
     expect(store.get(notebookClientAtom)).toBe(result.current);
-    act(() =>
-      fixtureMessages("notebook-events.sse").forEach(transport.receive),
-    );
+    const [connection] = fake.opened;
+    act(() => {
+      connection.open();
+      fixtureMessages("notebook-events.sse").forEach(connection.receive);
+    });
     expect(store.get(agentAttachmentsAtom)).toEqual([
       {
         id: "a2",
@@ -101,49 +176,48 @@ describe("useNotebookClient", () => {
   });
 
   it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
-    const first = makeTransport();
-    vi.mocked(createConnectionTransport).mockReturnValue(first.transport);
+    const fake = makeTransport();
+    vi.mocked(createNotebookTransport).mockReturnValue(fake.transport);
     const { rerender, options, unmount } = setup();
     await act(async () => Promise.resolve());
     const onNotification = vi.fn();
     rerender({ ...options, onNotification, url: () => "ws://redirected/ws" });
-    act(() => first.notify());
+    act(() => fake.notify());
     expect(options.onNotification).not.toHaveBeenCalled();
     expect(onNotification).toHaveBeenCalledExactlyOnceWith({
       data: { op: "reconnected" },
     });
-    expect(createConnectionTransport).toHaveBeenCalledOnce();
-    expect(vi.mocked(createConnectionTransport).mock.calls[0][0].url()).toBe(
+    expect(createNotebookTransport).toHaveBeenCalledOnce();
+    const [{ url }] = vi.mocked(createNotebookTransport).mock.calls[0];
+    expect(String(await (url as () => Promise<string>)())).toBe(
       "ws://redirected/ws",
     );
     unmount();
   });
 
-  it("cleans up every transport under Strict Mode and only connects the active lifetime", async () => {
-    const transports: ReturnType<typeof makeTransport>[] = [];
-    vi.mocked(createConnectionTransport).mockImplementation(() => {
+  it("closes every stream under Strict Mode and only follows the active lifetime", async () => {
+    const fakes: ReturnType<typeof makeTransport>[] = [];
+    vi.mocked(createNotebookTransport).mockImplementation(() => {
       const next = makeTransport();
-      transports.push(next);
+      fakes.push(next);
       return next.transport;
     });
     const { result, store, options, unmount } = setup(true);
     await act(async () => Promise.resolve());
-    expect(transports).toHaveLength(2);
-    expect(transports[0].transport.reconnect).not.toHaveBeenCalled();
-    expect(transports[0].transport.close).toHaveBeenCalledOnce();
-    expect(transports[1].transport.reconnect).toHaveBeenCalledOnce();
+    expect(fakes).toHaveLength(2);
+    expect(fakes[0].opened.every((attempt) => attempt.signal.aborted)).toBe(
+      true,
+    );
+    const [active] = fakes[1].opened;
+    act(() => active.open());
     expect(store.get(result.current.connectionAtom)).toEqual({
       state: ConnectionState.OPEN,
     });
-    act(() => {
-      for (const transport of transports) {
-        transport.notify();
-      }
-    });
+    act(() => fakes.forEach((fake) => fake.notify()));
     expect(options.onNotification).toHaveBeenCalledOnce();
     unmount();
-    expect(transports[1].transport.close).toHaveBeenCalledOnce();
+    expect(active.signal.aborted).toBe(true);
     await result.current.reconnect();
-    expect(transports[1].transport.reconnect).toHaveBeenCalledOnce();
+    expect(fakes[1].opened).toHaveLength(1);
   });
 });