@marimo-team/islands 0.25.3-dev10 → 0.25.3-dev12

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

@@ -0,0 +1,191 @@
+/* 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";
+
+function setup() {
+  const events = new EventTarget();
+  const healthy = new Deferred<void>();
+  const transport = {
+    readyState: WebSocket.CLOSED,
+    reconnect: vi.fn(),
+    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: () => transport,
+      url: () => "ws://localhost/ws",
+      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,
+  };
+}
+
+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.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();
+  });
+});

package/src/core/host/connection.ts

@@ -0,0 +1,81 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+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 ConnectionClosedReason = {
+  KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
+  KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR",
+} as const;
+
+export type ConnectionClosedReason =
+  (typeof ConnectionClosedReason)[keyof typeof ConnectionClosedReason];
+
+export type ConnectionStatus =
+  | {
+      state: typeof ConnectionState.CLOSED;
+      code: ConnectionClosedReason;
+      reason: string;
+    }
+  | {
+      state: typeof ConnectionState.CONNECTING;
+    }
+  | {
+      state:
+        | typeof ConnectionState.OPEN
+        | typeof ConnectionState.CLOSING
+        | typeof ConnectionState.NOT_STARTED;
+    };
+
+/** The recovery policy for a closed connection. */
+export interface CloseDecision {
+  kind: "terminal" | "gave-up" | "retry";
+  status: ConnectionStatus;
+}
+
+export type ConnectionEvent = "open" | "close" | "message" | "error";
+
+export interface IConnectionTransportMap {
+  open: WebSocketEventMap["open"];
+  message: WebSocketEventMap["message"];
+  close: WebSocketEventMap["close"];
+  error: WebSocketEventMap["error"];
+}
+
+export type ConnectionTransportCallback<T extends ConnectionEvent> = (
+  data: IConnectionTransportMap[T],
+) => void;
+
+/** Transport readiness, using the WebSocket readyState values. */
+export const TransportReadyState = {
+  CONNECTING: 0,
+  OPEN: 1,
+  CLOSING: 2,
+  CLOSED: 3,
+} as const;
+
+export type TransportReadyState =
+  (typeof TransportReadyState)[keyof typeof TransportReadyState];
+
+/** A reconnectable source of notebook messages. */
+export interface NotebookTransport {
+  reconnect(code?: number | undefined, reason?: string | undefined): void;
+  close(): void;
+  addEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void;
+  removeEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void;
+  readyState: TransportReadyState;
+}

package/src/core/host/events.ts

@@ -0,0 +1,60 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { z } from "zod";
+import * as Protocol from "./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.
+ */
+
+/** Event kinds requested by the client; removals and controls are implicit. */
+export const NOTEBOOK_EVENT_TYPES = [
+  "notebook",
+  "document",
+  "cell",
+  "marimo",
+] as const satisfies readonly (Protocol.ObjectType | "marimo")[];
+
+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 decoded event, or the name and cursor of an unrecognized event. */
+export type NotebookEvent =
+  | Protocol.EventOutput
+  | {
+      readonly event: "unknown";
+      readonly name: string;
+      readonly id?: string;
+    };
+
+const sessionEnvelope = z.object({ data: Protocol.MarimoNotification });
+
+/**
+ * Decode a host event or a legacy `{ data: notification }` session 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 (knownNames.has(parsed.event)) {
+    return Protocol.Event.parse(parsed);
+  }
+  return { event: "unknown", name: parsed.event, id: parsed.id };
+}

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

@@ -0,0 +1,226 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { Logger } from "@/utils/Logger";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import {
+  type ConnectionStatus,
+  ConnectionState,
+  ConnectionClosedReason,
+  type CloseDecision,
+  type NotebookTransport,
+  TransportReadyState,
+} from "./connection";
+import { parseNotebookEvent } from "./events";
+
+/** A connection to a notebook on a marimo host. */
+export interface NotebookClient {
+  /** Transport state, independent of kernel readiness. */
+  readonly connectionAtom: Atom<ConnectionStatus>;
+  /** Attempt to restore a closed connection after checking host health. */
+  reconnect(): Promise<void>;
+}
+
+/** The transport and application callbacks needed to connect 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;
+  /** Host events require an intact stream; legacy notifications are independent. */
+  messageFormat: "host" | "session";
+  /** Abort to close the transport 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;
+  /** Report a malformed message without coupling the client to editor UI. */
+  onMessageError: (error: unknown) => void;
+  /** Choose the connection status and retry policy after the transport closes. */
+  classifyClose: (event: CloseEvent) => CloseDecision;
+}
+
+/** Create a client without opening a connection. */
+export function createNotebookClient(store: ReturnType<typeof createStore>) {
+  const connectionAtom = atom<ConnectionStatus>({
+    state: ConnectionState.NOT_STARTED,
+  });
+  let reconnect: (() => Promise<void>) | undefined;
+  const client: NotebookClient = {
+    connectionAtom,
+    reconnect: async () => reconnect?.(),
+  };
+
+  function connect({
+    createTransport,
+    url,
+    messageFormat,
+    signal,
+    waitUntilReady,
+    checkHealth,
+    onNotification,
+    onMessageError,
+    decodeJson,
+    classifyClose,
+  }: NotebookConnection) {
+    if (signal.aborted) {
+      return;
+    }
+    let shouldRetry = true;
+    let terminal = false;
+    let checkingHealth = false;
+    const transport = createTransport(url);
+    const setStatus = (status: ConnectionStatus) =>
+      store.set(connectionAtom, status);
+    const disconnected = {
+      state: ConnectionState.CLOSED,
+      code: ConnectionClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    } satisfies ConnectionStatus;
+
+    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;
+      }
+    }
+
+    function onError(event: Event) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      Logger.warn("Notebook connection error", event);
+      setStatus(disconnected);
+      retry();
+    }
+
+    function onMessage(message: MessageEvent<string>) {
+      if (signal.aborted || terminal) {
+        return;
+      }
+      let event;
+      try {
+        event = parseNotebookEvent(decodeJson(message.data));
+      } catch (error) {
+        Logger.error("Failed to decode notebook message", error);
+        if (messageFormat === "session") {
+          onMessageError(error);
+          return;
+        }
+        terminal = true;
+        setStatus({
+          ...disconnected,
+          reason: "Failed to decode notebook message",
+        });
+        transport.close();
+        return;
+      }
+      if (event.event === "marimo") {
+        onNotification(event.data.marimo);
+      }
+    }
+
+    reconnect = async () => {
+      if (
+        signal.aborted ||
+        checkingHealth ||
+        transport.readyState === TransportReadyState.OPEN ||
+        transport.readyState === TransportReadyState.CONNECTING
+      ) {
+        return;
+      }
+      checkingHealth = true;
+      setStatus({ state: ConnectionState.CONNECTING });
+      try {
+        const healthy = await checkHealth();
+        if (signal.aborted) {
+          return;
+        }
+        if (!healthy) {
+          shouldRetry = false;
+          setStatus(disconnected);
+          return;
+        }
+        terminal = false;
+        shouldRetry = true;
+        transport.reconnect();
+      } catch (error) {
+        if (!signal.aborted) {
+          Logger.error("Notebook health check failed", error);
+          shouldRetry = false;
+          setStatus(disconnected);
+        }
+      } finally {
+        checkingHealth = false;
+      }
+    };
+
+    signal.addEventListener(
+      "abort",
+      () => {
+        transport.removeEventListener("open", onOpen);
+        transport.removeEventListener("close", onClose);
+        transport.removeEventListener("error", onError);
+        transport.removeEventListener("message", onMessage);
+        transport.close();
+        reconnect = undefined;
+        setStatus({ state: ConnectionState.NOT_STARTED });
+      },
+      { 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();
+          }
+        });
+    }
+  }
+
+  return { client, connect };
+}

package/src/core/run-app.tsx

@@ -22,7 +22,7 @@ import { getSessionId } from "./kernel/session";
 import { connectionAtom } from "./network/connection";
 import { useRequestClient } from "./network/requests";
 import { isAppConnecting } from "./websocket/connection-utils";
-import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
+import { useNotebookSession } from "./websocket/useNotebookSession";
 
 interface AppProps {
   appConfig: AppConfig;
@@ -51,7 +51,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
     };
   }, []);
 
-  const { connection, reconnect } = useMarimoKernelConnection({
+  const { connection, reconnect } = useNotebookSession({
     autoInstantiate: true,
     setCells: setCells,
     sessionId: getSessionId(),

package/src/core/runtime/__tests__/runtime.test.ts

@@ -2,14 +2,19 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
-import type { SessionId } from "@/core/kernel/session";
+import {
+  generateSessionId,
+  isSessionId,
+  type SessionId,
+} from "@/core/kernel/session";
 import { Logger } from "@/utils/Logger";
 import { RuntimeManager } from "../runtime";
 import type { RuntimeConfig } from "../types";
 
 // Mock the session module
-vi.mock("@/core/kernel/session", () => ({
-  getSessionId: () => "test-session-id" as SessionId,
+vi.mock("@/core/kernel/session", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/kernel/session")>()),
+  getSessionId: () => "test-session-id",
 }));
 
 vi.mock("@/utils/Logger", () => ({ Logger: Mocks.quietLogger() }));
@@ -51,6 +56,38 @@ describe("RuntimeManager", () => {
   });
 
   describe("getWsURL", () => {
+    it("uses host notebook events with the existing session and auth parameters", () => {
+      const original = window.location.href;
+      window.history.replaceState({}, "", "/?types=run");
+      try {
+        const runtime = new RuntimeManager({
+          ...mockConfig,
+          url: "https://example.com/nested?file=notebook.py",
+          notebookId: "nb_123",
+        });
+        const sessionId = "s_123456";
+        if (!isSessionId(sessionId)) {
+          throw new Error("Invalid test session ID");
+        }
+        const url = runtime.getWsURL(sessionId);
+        expect({
+          endpoint: `${url.origin}${url.pathname}`,
+          params: Object.fromEntries(url.searchParams),
+        }).toEqual({
+          endpoint:
+            "wss://example.com/nested/api/marimo/v1/notebooks/nb_123/events/ws",
+          params: {
+            file: "notebook.py",
+            types: "marimo",
+            session_id: sessionId,
+            access_token: "test-token",
+          },
+        });
+      } finally {
+        window.history.replaceState({}, "", original);
+      }
+    });
+
     it("should return WebSocket URL with session ID for https", () => {
       const runtime = new RuntimeManager(mockConfig);
       const sessionId = "1234" as SessionId;
@@ -59,6 +96,7 @@ describe("RuntimeManager", () => {
       expect(url.protocol).toBe("wss:");
       expect(url.hostname).toBe("example.com");
       expect(url.pathname).toBe("/ws");
+      expect(url.searchParams.has("types")).toBe(false);
       expect(url.searchParams.get("session_id")).toBe(sessionId);
     });
 
@@ -219,14 +257,20 @@ describe("RuntimeManager", () => {
   });
 
   describe("getSseURL", () => {
-    it("should return an http(s) URL with the session ID", () => {
-      const runtime = new RuntimeManager(mockConfig);
-      const url = runtime.getSseURL("1234" as SessionId);
-
-      expect(url.protocol).toBe("https:");
-      expect(url.pathname).toBe("/sse");
-      expect(url.searchParams.get("session_id")).toBe("1234");
-    });
+    it.each([undefined, "nb_123"])(
+      "uses the legacy session stream when notebookId is %s",
+      (notebookId) => {
+        const runtime = new RuntimeManager({
+          ...mockConfig,
+          url: "https://example.com/nested?file=notebook.py",
+          notebookId,
+        });
+        const sessionId = generateSessionId();
+        expect(runtime.getSseURL(sessionId).toString()).toBe(
+          `https://example.com/nested/sse?file=notebook.py&session_id=${sessionId}`,
+        );
+      },
+    );
 
     it("should never put the auth token in the URL, even cross-origin", () => {
       // SSE requests can send headers, so auth travels in the