@marimo-team/islands 0.25.3-dev80 → 0.25.3-dev82

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/useNotebookClient.ts

@@ -1,9 +1,9 @@
 /* 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 { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import type {
   NotificationMessage,
   NotificationPayload,
@@ -12,14 +12,16 @@ import {
   createNotebookClient,
   type NotebookClient,
 } from "../host/notebook-client";
-import { classifyCloseEvent } from "./connection-state";
 import { notebookClientAtom } from "./state";
 import {
-  createConnectionTransport,
-  type ConnectionTransportOptions,
-} from "./transports/create";
+  createNotebookTransport,
+  type NotebookTransportOptions,
+} from "./transports/notebook";
 
-type Options = ConnectionTransportOptions & {
+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;
@@ -37,24 +39,33 @@ export function useNotebookClient(options: Options): NotebookClient {
 
   useEffect(() => {
     const lifetime = new AbortController();
+    let needsInitialReadiness = true;
     store.set(notebookClientAtom, controller.client);
     controller.connect({
-      createTransport: (url) =>
-        createConnectionTransport({
-          url,
-          headers,
-          static: isStatic,
-          transportType,
-        }),
-      url,
+      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,
-      waitUntilReady,
       checkHealth,
       // The host validates the envelope, not the native notification contents.
       onNotification: (notification) =>
         onNotification({ data: notification as NotificationMessage }),
-      decodeJson: jsonParseWithSpecialChar,
-      classifyClose: classifyCloseEvent,
     });
     return () => {
       lifetime.abort();

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

@@ -1,49 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { describe, expect, expectTypeOf, it } from "vitest";
-import { z } from "zod";
-import { type NotebookEvent, 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 }],
-    };
-    expectTypeOf<
-      Extract<NotebookEvent, { event: "marimo" }>["data"]["marimo"]
-    >().toBeUnknown();
-    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/events.ts

@@ -1,68 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { z } from "zod";
-import * as Protocol from "@marimo-team/client/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.
- */
-
-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 native marimo message carried by a host stream.
- *
- * The protocol leaves extension events to their producer. The payload is
- * passed through for the session handler to interpret.
- */
-const MarimoEvent = z.object({
-  event: z.literal("marimo"),
-  id: z.string().min(1).optional(),
-  data: z.object({ marimo: z.unknown() }),
-});
-
-/** A decoded event, or the name and cursor of an unrecognized event. */
-export type NotebookEvent =
-  | Protocol.EventOutput
-  | z.output<typeof MarimoEvent>
-  | {
-      readonly event: "unknown";
-      readonly name: string;
-      readonly id?: string;
-    };
-
-const sessionEnvelope = z.object({ data: z.unknown() });
-
-/**
- * Decode a host event or an in-process WASM `{ data: notification }` 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 (parsed.event === "marimo") {
-    return MarimoEvent.parse(parsed);
-  }
-  if (knownNames.has(parsed.event)) {
-    return Protocol.Event.parse(parsed);
-  }
-  return { event: "unknown", name: parsed.event, id: parsed.id };
-}

package/src/core/websocket/__tests__/createConnectionTransport.test.ts

@@ -1,49 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { afterEach, describe, expect, it, vi } from "vitest";
-import { BasicTransport } from "../transports/basic";
-import { SseTransport } from "../transports/sse";
-import { WsTransport } from "../transports/ws";
-import { createConnectionTransport } from "../transports/create";
-
-vi.mock("../../wasm/utils", () => ({
-  isWasm: () => false,
-}));
-
-describe("createConnectionTransport", () => {
-  afterEach(() => {
-    vi.restoreAllMocks();
-  });
-
-  const options = {
-    url: () => "http://localhost/ws",
-    headers: () => ({}),
-  };
-
-  it("uses BasicTransport for static notebooks", () => {
-    const transport = createConnectionTransport({
-      ...options,
-      static: true,
-      transportType: "websocket",
-    });
-    expect(transport).toBeInstanceOf(BasicTransport);
-  });
-
-  it("uses WsTransport by default", () => {
-    const transport = createConnectionTransport({
-      ...options,
-      static: false,
-      transportType: "websocket",
-    });
-    expect(transport).toBeInstanceOf(WsTransport);
-  });
-
-  it("uses SseTransport when the transport type is sse", () => {
-    const transport = createConnectionTransport({
-      ...options,
-      static: false,
-      transportType: "sse",
-    });
-    expect(transport).toBeInstanceOf(SseTransport);
-  });
-});