@marimo-team/islands 0.25.3-dev141 → 0.25.3-dev143

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

@@ -0,0 +1,159 @@
+/* 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 { 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;
+  }
+}
+
+/** Send each request with the runtime's headers, read before each request. */
+export function withHeaders(
+  headers: () => Record<string, string>,
+): typeof globalThis.fetch {
+  return (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 });
+  };
+}
+
+/** 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: withHeaders(headers),
+      stop: (frame) =>
+        frame.event === "close"
+          ? stopOn(frame.data as { code?: number; reason?: string })
+          : undefined,
+    }),
+  );
+}
+
+/** A connection that opens and delivers nothing, for static notebooks. */
+export const idle: NotebookTransport = ({ signal }, { open }) =>
+  new Promise((settle) => {
+    open();
+    signal.addEventListener("abort", () => settle(), { once: true });
+  });

package/src/core/websocket/types.ts

@@ -1,47 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { ConnectionStatus as HostConnectionStatus } from "../host";
 import type { NotificationMessageData } from "../kernel/messages";
 
+export {
+  ConnectionState,
+  ConnectionClosedReason as WebSocketClosedReason,
+} from "../host";
+
 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/useNotebookReplica.ts

@@ -0,0 +1,69 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The editor session's replica, following its backend while mounted. */
+
+import { useStore } from "jotai";
+import { useEffect, useMemo } from "react";
+import useEvent from "react-use-event-hook";
+import type {
+  NotificationMessage,
+  NotificationPayload,
+} from "../kernel/messages";
+import { createNotebookReplica, type NotebookReplica } from "../host";
+import { notebookReplicaAtom } from "./state";
+import { createEditorBackend, type EditorBackendOptions } from "./backend";
+
+type Options = EditorBackendOptions & {
+  onNotification: (notification: NotificationPayload) => void;
+};
+
+export function useNotebookReplica(options: Options): NotebookReplica {
+  const store = useStore();
+  const { static: isStatic, transportType } = options;
+  const replica = useMemo(() => createNotebookReplica(store), [store]);
+  const url = useEvent(options.url);
+  const headers = useEvent(options.headers);
+  const notebookUrl = useEvent(options.notebookUrl);
+  const waitUntilReady = useEvent(options.waitUntilReady);
+  const checkHealth = useEvent(options.checkHealth);
+  const onNotification = useEvent(options.onNotification);
+  // Opens nothing until the replica follows it.
+  const backend = useMemo(
+    () =>
+      createEditorBackend({
+        static: isStatic,
+        transportType,
+        url,
+        headers,
+        notebookUrl,
+        waitUntilReady,
+        checkHealth,
+      }),
+    [
+      isStatic,
+      transportType,
+      url,
+      headers,
+      notebookUrl,
+      waitUntilReady,
+      checkHealth,
+    ],
+  );
+
+  useEffect(() => {
+    store.set(notebookReplicaAtom, replica);
+    const disconnect = replica.connect(backend, {
+      // The host validates the envelope, not the native notification contents.
+      onNotification: (notification) =>
+        onNotification({ data: notification as NotificationMessage }),
+    });
+    return () => {
+      disconnect();
+      if (store.get(notebookReplicaAtom) === replica) {
+        store.set(notebookReplicaAtom, null);
+      }
+    };
+  }, [replica, backend, store, onNotification]);
+
+  return replica;
+}

package/src/core/websocket/{useMarimoKernelConnection.tsx → useNotebookSession.tsx}

@@ -1,8 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { invalidatePackageData } from "@/core/packages/package-data";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
-import { useRef } from "react";
+import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
+import { useLayoutEffect } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
 import { getNotebook, useCellActions } from "@/core/cells/cells";
@@ -12,8 +12,7 @@ import type {
   NotificationMessageData,
   NotificationPayload,
 } from "@/core/kernel/messages";
-import { TRANSPORT_EXHAUSTED_REASON } from "@/core/websocket/transports/ws";
-import { useConnectionTransport } from "@/core/websocket/useWebSocket";
+import { useNotebookReplica } from "@/core/websocket/useNotebookReplica";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
@@ -21,11 +20,7 @@ import {
 } from "@/plugins/impl/anywidget/registry";
 import { logNever } from "@/utils/assertNever";
 import { prettyError } from "@/utils/errors";
-import {
-  type JsonString,
-  safeExtractSetUIElementMessageBuffers,
-} from "@/utils/json/base64";
-import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
 import { useAlertActions } from "../alerts/state";
@@ -80,105 +75,10 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
-import { isWasm } from "../wasm/utils";
-import {
-  type ConnectionStatus,
-  WebSocketClosedReason,
-  ConnectionState,
-} from "./types";
+import { WebSocketClosedReason, ConnectionState } from "./types";
 
 const SUPPORTS_LAZY_KERNELS = true;
 
-// All MARIMO_* reasons except TRANSPORT_EXHAUSTED are emitted by the backend
-// (marimo/_server/api/endpoints/ws_endpoint.py and ws/*.py). Keep in sync with
-// the backend literals. Over the SSE transport these arrive as in-band
-// `close` events rather than WebSocket close frames, so every terminal
-// reason must be classified here — an unclassified reason falls into the
-// retry path.
-export type CloseReason =
-  | "MARIMO_NO_FILE_KEY"
-  | "MARIMO_NO_SESSION_ID"
-  | "MARIMO_NO_SESSION"
-  | "MARIMO_SHUTDOWN"
-  | "MARIMO_KERNEL_STARTUP_ERROR"
-  | "MARIMO_UNAUTHORIZED"
-  | "MARIMO_KIOSK_NOT_ALLOWED"
-  | typeof TRANSPORT_EXHAUSTED_REASON;
-
-export type CloseDecision =
-  | { kind: "terminal"; status: ConnectionStatus; closeTransport: boolean }
-  | { kind: "gave-up"; status: ConnectionStatus }
-  | { kind: "retry"; status: ConnectionStatus };
-
-export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
-  switch (event.reason as CloseReason | undefined) {
-    case TRANSPORT_EXHAUSTED_REASON:
-      return {
-        kind: "gave-up",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-      };
-    case "MARIMO_NO_FILE_KEY":
-    case "MARIMO_NO_SESSION_ID":
-    case "MARIMO_NO_SESSION":
-    case "MARIMO_SHUTDOWN":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_UNAUTHORIZED":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "not authorized",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_KIOSK_NOT_ALLOWED":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kiosk mode is not available for this session",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_KERNEL_STARTUP_ERROR":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
-          reason: "Failed to start kernel sandbox",
-        },
-        closeTransport: true,
-      };
-    default:
-      // Empty/undefined reasons are normal transient closes. Anything else is
-      // an unknown server reason; warn so a new MARIMO_* reason doesn't fall
-      // silently into the retry path.
-      if (event.reason) {
-        logNever(event.reason as never);
-      }
-  }
-
-  return {
-    kind: "retry",
-    status: { state: ConnectionState.CONNECTING },
-  };
-}
-
 function getExistingCells(): CellData[] | undefined {
   if (!SUPPORTS_LAZY_KERNELS) {
     return undefined;
@@ -190,16 +90,12 @@ function getExistingCells(): CellData[] | undefined {
   );
 }
 
-/**
- * Creates a connection to the Marimo kernel and handles incoming messages.
- */
-export function useMarimoKernelConnection(opts: {
+/** Coordinate session startup and apply notifications to the editor. */
+export function useNotebookSession(opts: {
   sessionId: SessionId;
   autoInstantiate: boolean;
   setCells: (cells: CellData[], layout: LayoutState) => void;
 }) {
-  // Track whether we want to try reconnecting.
-  const shouldTryReconnecting = useRef<boolean>(true);
   const { autoInstantiate, sessionId, setCells } = opts;
   const { showBoundary } = useErrorBoundary();
 
@@ -246,8 +142,7 @@ export function useMarimoKernelConnection(opts: {
     filterFromVariables: filterStorageFromVariables,
   } = useStorageActions();
 
-  const handleMessage = (e: MessageEvent<JsonString<NotificationPayload>>) => {
-    const msg = jsonParseWithSpecialChar(e.data);
+  const handleNotification = (msg: NotificationPayload) => {
     switch (msg.data.op) {
       case "reload":
         reloadSafe();
@@ -368,13 +263,18 @@ export function useMarimoKernelConnection(opts: {
         filterStorageFromVariables(variableNames);
         return;
       case "variable-values":
-        setMetadata(
-          msg.data.variables.map((v) => ({
+        setMetadata([
+          ...(msg.data.removed ?? []).map((v) => ({
+            name: v.name as VariableName,
+            dataType: undefined,
+            value: undefined,
+          })),
+          ...msg.data.variables.map((v) => ({
             name: v.name as VariableName,
             dataType: v.datatype,
             value: v.value,
           })),
-        );
+        ]);
         return;
       case "alert":
         toast({
@@ -505,107 +405,56 @@ export function useMarimoKernelConnection(opts: {
     }
   };
 
-  const tryReconnecting = (code?: number, reason?: string) => {
-    // If not properly gated, we could try reconnecting forever if the
-    // issue is not transient. So we want to try reconnecting only once after an
-    // open connection is closed.
-    if (shouldTryReconnecting.current) {
-      shouldTryReconnecting.current = false;
-      ws.reconnect(code, reason);
-    }
-  };
-
-  // Manual reconnect. Probes /health first to fail fast when the runtime
-  // is unreachable, instead of waiting on partysocket's retry budget.
-  const reconnect = async () => {
-    if (
-      ws.readyState === WebSocket.OPEN ||
-      ws.readyState === WebSocket.CONNECTING
-    ) {
-      return;
-    }
-    shouldTryReconnecting.current = true;
-    setKernelStartupError(null);
-    setConnection({ state: ConnectionState.CONNECTING });
-    const healthy = await runtimeManager.reconcileFromHealth();
-    if (!healthy) {
-      shouldTryReconnecting.current = false;
-      setConnection({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-      return;
-    }
-    ws.reconnect();
+  const store = useStore();
+  const onMessageError = (error: unknown) => {
+    toast({
+      title: "Failed to handle message",
+      description: prettyError(error),
+      variant: "danger",
+    });
   };
-
-  const ws = useConnectionTransport({
+  const replica = useNotebookReplica({
     static: isStaticNotebook(),
     transportType,
-    /**
-     * Unique URL for this session.
-     */
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
         : runtimeManager.getWsURL(sessionId).toString(),
-    /**
-     * Auth headers for the SSE transport. WebSockets cannot send headers
-     * and instead put the access token in the URL.
-     */
+    // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
-
-    /**
-     * The transport is open; kernel-ready establishes session readiness.
-     */
-    onOpen: async () => {
-      updateStartupProgress(null);
-      // If we are open, we can reset our reconnecting flag.
-      shouldTryReconnecting.current = true;
-    },
-
-    /**
-     * Wait to connect, in case the remote kernel still starting up.
-     */
-    waitToConnect: async () => {
-      if (isStaticNotebook()) {
-        return;
-      }
-      if (isWasm()) {
-        return;
-      }
-
-      if (runtimeManager.isSameOrigin) {
-        return;
+    notebookUrl: (notebookId) => runtimeManager.getNotebookURL(notebookId),
+    // Only the host's transports read the URL, so static and WASM never wait.
+    waitUntilReady: async () => {
+      if (!runtimeManager.isSameOrigin) {
+        await runtimeManager.waitForHealthy();
       }
-      await runtimeManager.waitForHealthy();
     },
-
-    /**
-     * Handle messages sent by the kernel.
-     */
-    onMessage: (e) => {
+    checkHealth: () => runtimeManager.reconcileFromHealth(),
+    onNotification: (notification) => {
       try {
-        handleMessage(e);
+        handleNotification(notification);
       } catch (error) {
-        Logger.error("Failed to handle message", e.data, error);
-        toast({
-          title: "Failed to handle message",
-          description: prettyError(error),
-          variant: "danger",
-        });
+        Logger.error("Failed to handle notification", notification, error);
+        onMessageError(error);
       }
     },
+  });
 
-    /**
-     * Handle a close event. We may want to reconnect.
-     */
-    onClose: (e) => {
-      Logger.warn("WebSocket closed", e.code, e.reason);
-      const decision = classifyCloseEvent(e);
+  // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
+  useLayoutEffect(() => {
+    const updateConnection = () => {
+      const status = store.get(replica.connectionAtom);
+      if (status.state === ConnectionState.NOT_STARTED) {
+        return;
+      }
+      if (status.state === ConnectionState.OPEN) {
+        updateStartupProgress(null);
+        return; // kernel-ready, not the socket, establishes session readiness.
+      }
+      if (status.state === ConnectionState.CONNECTING) {
+        setKernelStartupError(null);
+      }
       setConnection((previous) => {
-        const status = decision.status;
         if (
           status.state === ConnectionState.CLOSED &&
           status.code === WebSocketClosedReason.KERNEL_STARTUP_ERROR &&
@@ -624,31 +473,17 @@ export function useMarimoKernelConnection(opts: {
         }
         return status;
       });
-      if (decision.kind === "terminal" && decision.closeTransport) {
-        ws.close(); // close to prevent reconnecting
-        return;
-      }
-      if (decision.kind === "retry") {
-        // Session should be valid
-        // - browser tab might have been closed or re-opened
-        // - computer might have just woken from sleep
-        tryReconnecting(e.code, e.reason);
-      }
-    },
-
-    /**
-     * When we encounter an error, we should close the connection.
-     */
-    onError: (e) => {
-      Logger.warn("WebSocket error", e);
-      setConnection({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-      tryReconnecting();
-    },
-  });
-
-  return { connection, reconnect };
+    };
+    const unsubscribe = store.sub(replica.connectionAtom, updateConnection);
+    updateConnection();
+    return unsubscribe;
+  }, [
+    replica,
+    store,
+    setConnection,
+    setKernelStartupError,
+    updateStartupProgress,
+  ]);
+
+  return { replica, connection, reconnect: replica.reconnect };
 }

package/src/css/globals.css

@@ -75,6 +75,9 @@
     --ring: hsl(215deg 20.2% 65.1%);
 
     /* Semantic colors. */
+    --brand: light-dark(#365f43, #a8c9ae);
+    --brand-hover: light-dark(#2b5037, #bdd8c2);
+    --brand-foreground: light-dark(#fff, #172b1c);
     --destructive: var(--red-9, #e5484d);
     --destructive-hover: var(--red-10, #dc3e42);
     --destructive-border: var(--red-11, #ce2c31);
@@ -89,6 +92,7 @@
     --action-hover: var(--yellow-10, #ffdc00);
     --action-border: var(--yellow-11, #9e6c00);
     --action-foreground: var(--yellow-12, #473b1f);
+    --search-highlight: var(--yellow-4, #fff394);
     --link: light-dark(hsl(211deg 90% 42%), hsl(211deg 90% 62%));
     --link-visited: light-dark(hsl(272deg 51% 54%), hsl(272deg 51% 74%));
     --stale: hsl(42deg 56% 44% / 25%);
@@ -137,6 +141,7 @@
     --action-hover: var(--yellow-7, #665417);
     --action-border: var(--yellow-11, #f5e147);
     --action-foreground: var(--yellow-12, #f6eeb4);
+    --search-highlight: var(--yellow-6, #524202);
     --stale: hsl(54deg 100% 86.7% / 25%);
   }
 }