@marimo-team/frontend 0.25.2 → 0.25.3-dev100

package/src/core/websocket/transports/notebook.ts

@@ -0,0 +1,213 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  type NotebookTransport,
+  sse,
+  StreamStopped,
+  type StreamUrl,
+} from "@marimo-team/client";
+import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { PyodideBridge } from "../../wasm/bridge";
+import { isWasm } from "../../wasm/utils";
+import { classifyCloseEvent } from "../connection-state";
+
+/** The server can be slow to accept when many notebooks are open. */
+const CONNECT_MS = 10_000;
+
+/**
+ * The error that ends a stream on this close, or none to retry.
+ *
+ * `cause` is the connection status the editor shows.
+ */
+function stopOn(close: { code?: number; reason?: string }) {
+  const decision = classifyCloseEvent(close);
+  if (decision.kind === "retry") {
+    return undefined;
+  }
+  const reason = "reason" in decision.status ? decision.status.reason : "";
+  return new StreamStopped(reason, { cause: decision.status });
+}
+
+/** Read a URL provider, stopping the stream when it gives up. */
+async function resolve(url: StreamUrl): Promise<URL | undefined> {
+  try {
+    return new URL(typeof url === "function" ? await url() : url);
+  } catch (error) {
+    if (error instanceof StreamStopped) {
+      throw error;
+    }
+    return undefined;
+  }
+}
+
+/** A resume cursor and replica belong to the session that produced them. */
+function withSessionIdentity(
+  url: StreamUrl,
+  connect: (target: URL) => NotebookTransport,
+): NotebookTransport {
+  let previous: string | undefined;
+  return async (attempt, handlers) => {
+    const target = await resolve(url);
+    if (target === undefined || attempt.signal.aborted) {
+      return;
+    }
+    target.searchParams.delete("cursor");
+    const identity = new URL(target);
+    identity.searchParams.delete("access_token");
+    identity.searchParams.sort();
+    identity.hash = "";
+    const key = identity.toString();
+    let { cursor } = attempt;
+    if (previous !== undefined && previous !== key) {
+      // Reset both the stream's cursor and its subscribers' state.
+      handlers.message({ event: "reset", data: {} });
+      cursor = null;
+    }
+    previous = key;
+    if (!attempt.signal.aborted) {
+      return connect(target)({ ...attempt, cursor }, handlers);
+    }
+  };
+}
+
+/**
+ * Follow a marimo session over its WebSocket route.
+ *
+ * Each attempt opens one socket with the cursor as a query parameter. The
+ * editor's close policy decides which closes end the stream.
+ *
+ * TODO: Notice a connection that died without closing, such as after a laptop
+ * sleeps. `/sse` sends a heartbeat that its silence limit relies on, and
+ * `/ws` drops it. Sending it on `/ws` would let this transport close a
+ * silent socket.
+ */
+export function websocket(url: StreamUrl): NotebookTransport {
+  return withSessionIdentity(
+    url,
+    (target) =>
+      async ({ cursor, signal }, { open, message }) => {
+        if (cursor !== null) {
+          target.searchParams.set("cursor", cursor);
+        }
+        await new Promise<void>((settle, reject) => {
+          const socket = new WebSocket(target);
+          const timeout = setTimeout(() => socket.close(), CONNECT_MS);
+          const abort = () => socket.close();
+          signal.addEventListener("abort", abort, { once: true });
+          socket.addEventListener("open", () => {
+            clearTimeout(timeout);
+            open();
+          });
+          socket.addEventListener("message", ({ data }) =>
+            // marimo's parser reads `NaN`. A message it cannot read arrives as an
+            // empty object, which `NotebookEvents` rejects as not an event.
+            message(jsonParseWithSpecialChar(String(data))),
+          );
+          socket.addEventListener("close", (event) => {
+            clearTimeout(timeout);
+            signal.removeEventListener("abort", abort);
+            const stopped = signal.aborted ? undefined : stopOn(event);
+            if (stopped === undefined) {
+              settle();
+            } else {
+              reject(stopped);
+            }
+          });
+        });
+      },
+  );
+}
+
+/**
+ * Follow a marimo session over its server-sent events route.
+ *
+ * The session's `close` frames carry the same reasons as WebSocket closes.
+ * marimo's parser decodes tagged bigints and non-finite numbers in every frame.
+ */
+export function sessionSse(
+  url: StreamUrl,
+  headers: () => Record<string, string>,
+): NotebookTransport {
+  return withSessionIdentity(url, (target) =>
+    sse(target, {
+      parseData: jsonParseWithSpecialChar,
+      fetch: (input, init) => {
+        const sent = new Headers(init?.headers);
+        for (const [name, value] of Object.entries(headers())) {
+          sent.set(name, value);
+        }
+        return globalThis.fetch(input, { ...init, headers: sent });
+      },
+      stop: (frame) =>
+        frame.event === "close"
+          ? stopOn(frame.data as { code?: number; reason?: string })
+          : undefined,
+    }),
+  );
+}
+
+/** A connection that opens and delivers nothing, for static notebooks. */
+const idle: NotebookTransport = ({ signal }, { open }) =>
+  new Promise((settle) => {
+    open();
+    signal.addEventListener("abort", () => settle(), { once: true });
+  });
+
+/**
+ * Follow the notebook running in the WASM worker.
+ *
+ * The worker sends marimo notifications as `{ data: notification }`, which
+ * arrive as `marimo` events.
+ */
+export function pyodide(
+  attach: (consumer: (message: MessageEvent<string>) => void) => void,
+): NotebookTransport {
+  return ({ signal }, { open, message }) =>
+    new Promise((settle) => {
+      open();
+      attach(({ data }) => {
+        if (signal.aborted) {
+          return;
+        }
+        const value = jsonParseWithSpecialChar(data);
+        // The worker sends notifications without the host's event envelope.
+        message(
+          typeof value === "object" &&
+            value !== null &&
+            !("event" in value) &&
+            "data" in value
+            ? { event: "marimo", data: { marimo: value.data } }
+            : value,
+        );
+      });
+      signal.addEventListener("abort", () => settle(), { once: true });
+    });
+}
+
+/** How the editor follows its notebook. */
+export interface NotebookTransportOptions {
+  static: boolean;
+  transportType: "websocket" | "sse";
+  /** The session route, read before each connection. */
+  url: StreamUrl;
+  /** The runtime's headers, read before each request. */
+  headers: () => Record<string, string>;
+}
+
+/** Choose the transport for the editor's notebook. */
+export function createNotebookTransport(
+  options: NotebookTransportOptions,
+): NotebookTransport {
+  if (options.static) {
+    return idle;
+  }
+  if (isWasm()) {
+    return pyodide((consumer) =>
+      PyodideBridge.INSTANCE.attachMessageConsumer(consumer),
+    );
+  }
+  if (options.transportType === "sse") {
+    return sessionSse(options.url, options.headers);
+  }
+  return websocket(options.url);
+}

package/src/core/websocket/types.ts

@@ -1,47 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { ConnectionStatus as HostConnectionStatus } from "../host/connection";
 import type { NotificationMessageData } from "../kernel/messages";
 
+export {
+  ConnectionState,
+  ConnectionClosedReason as WebSocketClosedReason,
+} from "../host/connection";
+
 export type ConnectionPhase =
   | NotificationMessageData<"startup-progress">["phase"]
   | "reconnecting";
 
-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 WebSocketClosedReason = {
-  KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
-  KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR",
-} as const;
-
-export type WebSocketClosedReason =
-  (typeof WebSocketClosedReason)[keyof typeof WebSocketClosedReason];
-
-export type ConnectionStatus =
-  | {
-      state: typeof ConnectionState.CLOSED;
-      code: WebSocketClosedReason;
-      /**
-       * Human-readable reason for closing the connection.
-       */
-      reason: string;
-      phase?: ConnectionPhase;
-    }
-  | {
-      state: typeof ConnectionState.CONNECTING;
-      phase?: ConnectionPhase;
-    }
-  | {
-      state:
-        | typeof ConnectionState.OPEN
-        | typeof ConnectionState.CLOSING
-        | typeof ConnectionState.NOT_STARTED;
-    };
+export type ConnectionStatus = HostConnectionStatus & {
+  phase?: ConnectionPhase;
+};

package/src/core/websocket/useNotebookClient.ts

@@ -0,0 +1,89 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { StreamStopped } from "@marimo-team/client";
+import { useStore } from "jotai";
+import { useEffect, useMemo } from "react";
+import useEvent from "react-use-event-hook";
+import type {
+  NotificationMessage,
+  NotificationPayload,
+} from "../kernel/messages";
+import {
+  createNotebookClient,
+  type NotebookClient,
+} from "../host/notebook-client";
+import { notebookClientAtom } from "./state";
+import {
+  createNotebookTransport,
+  type NotebookTransportOptions,
+} from "./transports/notebook";
+
+type Options = Omit<NotebookTransportOptions, "url"> & {
+  /** The session route, read before each connection. */
+  url: () => string;
+  /** Wait until the host can accept a connection. */
+  waitUntilReady: () => Promise<void>;
+  checkHealth: () => Promise<boolean>;
+  onNotification: (notification: NotificationPayload) => void;
+};
+
+export function useNotebookClient(options: Options): NotebookClient {
+  const store = useStore();
+  const { static: isStatic, transportType } = options;
+  const controller = useMemo(() => createNotebookClient(store), [store]);
+  const url = useEvent(options.url);
+  const headers = useEvent(options.headers);
+  const waitUntilReady = useEvent(options.waitUntilReady);
+  const checkHealth = useEvent(options.checkHealth);
+  const onNotification = useEvent(options.onNotification);
+
+  useEffect(() => {
+    const lifetime = new AbortController();
+    let needsInitialReadiness = true;
+    store.set(notebookClientAtom, controller.client);
+    controller.connect({
+      transport: createNotebookTransport({
+        static: isStatic,
+        transportType,
+        url: async () => {
+          if (needsInitialReadiness) {
+            // Explicit retries check current health, even if startup failed.
+            needsInitialReadiness = false;
+            try {
+              await waitUntilReady();
+            } catch (error) {
+              throw new StreamStopped("Healthy connection never made", {
+                cause: error,
+              });
+            }
+          }
+          return url();
+        },
+        headers,
+      }),
+      signal: lifetime.signal,
+      checkHealth,
+      // The host validates the envelope, not the native notification contents.
+      onNotification: (notification) =>
+        onNotification({ data: notification as NotificationMessage }),
+    });
+    return () => {
+      lifetime.abort();
+      if (store.get(notebookClientAtom) === controller.client) {
+        store.set(notebookClientAtom, null);
+      }
+    };
+  }, [
+    controller,
+    store,
+    isStatic,
+    transportType,
+    url,
+    headers,
+    waitUntilReady,
+    checkHealth,
+    onNotification,
+  ]);
+
+  return controller.client;
+}