@marimo-team/frontend 0.25.3-dev13 → 0.25.3-dev135

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

@@ -0,0 +1,190 @@
+/* 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 });
+  });
+
+/**
+ * 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 });
+    });
+}

package/src/core/websocket/types.ts

@@ -1,12 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { ConnectionStatus as HostConnectionStatus } from "../host/connection";
+import type { ConnectionStatus as HostConnectionStatus } from "../host";
 import type { NotificationMessageData } from "../kernel/messages";
 
 export {
   ConnectionState,
   ConnectionClosedReason as WebSocketClosedReason,
-} from "../host/connection";
+} from "../host";
 
 export type ConnectionPhase =
   | NotificationMessageData<"startup-progress">["phase"]

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/useNotebookSession.tsx

@@ -5,7 +5,6 @@ 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 { attachmentsAtom } from "@/core/attachments/state";
 import { getNotebook, useCellActions } from "@/core/cells/cells";
 import { applyTransactionChanges } from "@/core/cells/document-changes";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
@@ -13,7 +12,7 @@ import type {
   NotificationMessageData,
   NotificationPayload,
 } from "@/core/kernel/messages";
-import { useNotebookClient } from "@/core/websocket/useNotebookClient";
+import { useNotebookReplica } from "@/core/websocket/useNotebookReplica";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
@@ -66,7 +65,7 @@ import { type LayoutState, useLayoutActions } from "../layout/state";
 import { kioskModeAtom } from "../mode";
 import { connectionAtom, startupProgressAtom } from "../network/connection";
 import type { RequestId } from "../network/DeferredRequestRegistry";
-import { runtimeConfigAtom, useRuntimeManager } from "../runtime/config";
+import { useRuntimeManager } from "../runtime/config";
 import { SECRETS_REGISTRY } from "../secrets/request-registry";
 import { isStaticNotebook } from "../static/static-state";
 import {
@@ -76,7 +75,6 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
-import { isWasm } from "../wasm/utils";
 import { WebSocketClosedReason, ConnectionState } from "./types";
 
 const SUPPORTS_LAZY_KERNELS = true;
@@ -136,11 +134,9 @@ export function useNotebookSession(opts: {
   const setStableSessionId = useSetAtom(stableSessionIdAtom);
   const setCapabilities = useSetAtom(capabilitiesAtom);
   const runtimeManager = useRuntimeManager();
-  const { notebookId } = useAtomValue(runtimeConfigAtom);
   const transportType = useAtomValue(connectionTransportTypeAtom);
   const setCacheInfo = useSetAtom(cacheInfoAtom);
   const setKernelStartupError = useSetAtom(kernelStartupErrorAtom);
-  const setAttachments = useSetAtom(attachmentsAtom);
   const {
     setNamespaces: setStorageNamespaces,
     filterFromVariables: filterStorageFromVariables,
@@ -148,9 +144,6 @@ export function useNotebookSession(opts: {
 
   const handleNotification = (msg: NotificationPayload) => {
     switch (msg.data.op) {
-      case "attachments":
-        setAttachments(msg.data.attachments);
-        return;
       case "reload":
         reloadSafe();
         return;
@@ -185,7 +178,6 @@ export function useNotebookSession(opts: {
           existingCells,
         });
         if (!msg.data.resumed) {
-          setAttachments([]);
           // A freshly started kernel may expose a different environment
           // (e.g. new env vars); re-run discovery instead of reusing stale
           // suggestions from the previous kernel instance.
@@ -416,25 +408,23 @@ export function useNotebookSession(opts: {
       variant: "danger",
     });
   };
-  const notebook = useNotebookClient({
-    sessionId,
+  const replica = useNotebookReplica({
     static: isStaticNotebook(),
     transportType,
-    messageFormat:
-      notebookId && transportType === "websocket" ? "host" : "session",
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
         : runtimeManager.getWsURL(sessionId).toString(),
     // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
+    notebookUrl: (notebookId) => runtimeManager.getNotebookURL(notebookId),
+    // Only the host's transports read the URL, so static and WASM never wait.
     waitUntilReady: async () => {
-      if (!isStaticNotebook() && !isWasm() && !runtimeManager.isSameOrigin) {
+      if (!runtimeManager.isSameOrigin) {
         await runtimeManager.waitForHealthy();
       }
     },
     checkHealth: () => runtimeManager.reconcileFromHealth(),
-    onMessageError,
     onNotification: (notification) => {
       try {
         handleNotification(notification);
@@ -448,11 +438,10 @@ export function useNotebookSession(opts: {
   // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
   useLayoutEffect(() => {
     const updateConnection = () => {
-      const status = store.get(notebook.connectionAtom);
+      const status = store.get(replica.connectionAtom);
       if (status.state === ConnectionState.NOT_STARTED) {
         return;
       }
-      setAttachments([]);
       if (status.state === ConnectionState.OPEN) {
         updateStartupProgress(null);
         return; // kernel-ready, not the socket, establishes session readiness.
@@ -480,17 +469,16 @@ export function useNotebookSession(opts: {
         return status;
       });
     };
-    const unsubscribe = store.sub(notebook.connectionAtom, updateConnection);
+    const unsubscribe = store.sub(replica.connectionAtom, updateConnection);
     updateConnection();
     return unsubscribe;
   }, [
-    notebook,
+    replica,
     store,
     setConnection,
-    setAttachments,
     setKernelStartupError,
     updateStartupProgress,
   ]);
 
-  return { notebook, connection, reconnect: notebook.reconnect };
+  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%);
   }
 }

package/src/hooks/__tests__/useAsyncData.test.tsx

@@ -0,0 +1,148 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { act, render, renderHook, waitFor } from "@testing-library/react";
+import { Activity } from "react";
+import { describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { useAsyncData } from "../useAsyncData";
+
+function AsyncDataStatus({ fetch }: { fetch: () => Promise<string> }) {
+  const result = useAsyncData(fetch, []);
+  return (
+    <output
+      data-status={result.status}
+      data-fetching={result.isFetching}
+      data-error={result.error?.message ?? ""}
+    >
+      {result.data}
+    </output>
+  );
+}
+
+describe("useAsyncData", () => {
+  it.each(["success", "error"] as const)(
+    "restarts fetching when Activity reopens after %s",
+    async (initialStatus) => {
+      const initial = new Deferred<string>();
+      const restarted = new Deferred<string>();
+      const fetch = vi
+        .fn(() => restarted.promise)
+        .mockImplementationOnce(() => initial.promise);
+      const view = render(
+        <Activity mode="visible">
+          <AsyncDataStatus fetch={fetch} />
+        </Activity>,
+      );
+      await act(async () => {
+        if (initialStatus === "success") {
+          initial.resolve("first");
+        } else {
+          initial.reject(new Error("offline"));
+        }
+      });
+      expect(view.getByRole("status").dataset.status).toBe(initialStatus);
+
+      view.rerender(
+        <Activity mode="hidden">
+          <AsyncDataStatus fetch={fetch} />
+        </Activity>,
+      );
+      view.rerender(
+        <Activity mode="visible">
+          <AsyncDataStatus fetch={fetch} />
+        </Activity>,
+      );
+
+      expect(fetch).toHaveBeenCalledTimes(2);
+      expect(view.getByRole("status").dataset).toMatchObject({
+        status: initialStatus === "success" ? "loading" : "pending",
+        error: "",
+        fetching: "true",
+      });
+      expect(view.getByRole("status").textContent).toBe(
+        initialStatus === "success" ? "first" : "",
+      );
+      await act(async () => restarted.resolve("latest"));
+      expect(view.getByRole("status").dataset).toMatchObject({
+        status: "success",
+        fetching: "false",
+      });
+      expect(view.getByRole("status").textContent).toBe("latest");
+    },
+  );
+
+  it("ignores a request canceled when Activity hides", async () => {
+    const canceled = new Deferred<string>();
+    const restarted = new Deferred<string>();
+    const fetch = vi
+      .fn(() => restarted.promise)
+      .mockImplementationOnce(() => canceled.promise);
+    const view = render(
+      <Activity mode="visible">
+        <AsyncDataStatus fetch={fetch} />
+      </Activity>,
+    );
+    view.rerender(
+      <Activity mode="hidden">
+        <AsyncDataStatus fetch={fetch} />
+      </Activity>,
+    );
+    view.rerender(
+      <Activity mode="visible">
+        <AsyncDataStatus fetch={fetch} />
+      </Activity>,
+    );
+    expect(fetch).toHaveBeenCalledTimes(2);
+    await act(async () => canceled.resolve("obsolete"));
+    expect(view.getByRole("status").dataset.status).toBe("pending");
+    expect(view.getByRole("status").textContent).toBe("");
+    await act(async () => restarted.resolve("latest"));
+    expect(view.getByRole("status").dataset.status).toBe("success");
+    expect(view.getByRole("status").textContent).toBe("latest");
+  });
+
+  it("preserves data while refetching and ignores obsolete requests", async () => {
+    const requests = [
+      new Deferred<string>(),
+      new Deferred<string>(),
+      new Deferred<string>(),
+    ];
+    const fetch = vi
+      .fn(() => requests[0].promise)
+      .mockImplementationOnce(() => requests[0].promise)
+      .mockImplementationOnce(() => requests[1].promise)
+      .mockImplementationOnce(() => requests[2].promise);
+    const { result, rerender } = renderHook(
+      ({ key }) => useAsyncData(fetch, [key]),
+      { initialProps: { key: 0 } },
+    );
+    expect(result.current.status).toBe("pending");
+    await act(async () => requests[0].resolve("first"));
+    expect(result.current.data).toBe("first");
+    rerender({ key: 1 });
+    expect(result.current).toMatchObject({ status: "loading", data: "first" });
+    rerender({ key: 2 });
+    await act(async () => requests[1].resolve("obsolete"));
+    expect(result.current).toMatchObject({ status: "loading", data: "first" });
+    await act(async () => requests[2].resolve("latest"));
+    expect(result.current).toMatchObject({ status: "success", data: "latest" });
+  });
+
+  it("retries after errors and allows manual updates during a request", async () => {
+    const request = new Deferred<string>();
+    const fetch = vi
+      .fn<() => Promise<string>>()
+      .mockRejectedValueOnce(new Error("offline"))
+      .mockImplementation(() => request.promise);
+    const { result } = renderHook(() => useAsyncData(fetch, []));
+    await waitFor(() => expect(result.current.status).toBe("error"));
+    act(() => result.current.refetch());
+    expect(result.current.status).toBe("pending");
+    act(() => result.current.setData("manual"));
+    expect(result.current).toMatchObject({ status: "success", data: "manual" });
+    await act(async () => request.resolve("response"));
+    expect(result.current).toMatchObject({
+      status: "success",
+      data: "response",
+    });
+  });
+});

package/src/hooks/__tests__/useDebounce.test.tsx

@@ -121,6 +121,57 @@ describe("useDebounceControlledState", () => {
 
     expect(onChange).not.toHaveBeenCalled();
   });
+
+  it("replaces an uncommitted draft when the external value changes", () => {
+    const onChange = vi.fn();
+    const { result, rerender } = renderHook(
+      ({ value }) =>
+        useDebounceControlledState({
+          initialValue: value,
+          onChange,
+          delay: 1000,
+        }),
+      { initialProps: { value: "initial" } },
+    );
+
+    act(() => result.current.onChange("draft"));
+    rerender({ value: "external" });
+    expect(result.current.value).toBe("external");
+    act(() => vi.runAllTimers());
+    expect(result.current.debouncedValue).toBe("external");
+    expect(onChange).not.toHaveBeenCalled();
+  });
+
+  it("handles NaN external values without repeatedly resetting the draft", () => {
+    const onChange = vi.fn();
+    const { result, rerender } = renderHook(
+      ({ value }) =>
+        useDebounceControlledState({ initialValue: value, onChange }),
+      { initialProps: { value: Number.NaN } },
+    );
+
+    expect(result.current.value).toBeNaN();
+    act(() => result.current.onChange(1));
+    rerender({ value: Number.NaN });
+    expect(result.current.value).toBe(1);
+    rerender({ value: 2 });
+    expect(result.current.value).toBe(2);
+    rerender({ value: Number.NaN });
+    expect(result.current.value).toBeNaN();
+  });
+
+  it("resets the draft when the external value changes from zero to negative zero", () => {
+    const onChange = vi.fn();
+    const { result, rerender } = renderHook(
+      ({ value }) =>
+        useDebounceControlledState({ initialValue: value, onChange }),
+      { initialProps: { value: 0 } },
+    );
+
+    act(() => result.current.onChange(1));
+    rerender({ value: -0 });
+    expect(result.current.value).toBe(-0);
+  });
 });
 
 describe("useDebouncedCallback", () => {

package/src/hooks/__tests__/useDeepCompareMemoize.test.tsx

@@ -0,0 +1,39 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { useDeepCompareMemoize } from "../useDeepCompareMemoize";
+
+describe("useDeepCompareMemoize", () => {
+  it("keeps equal values stable and updates changed values immediately", () => {
+    const initial = { nested: [1, 2] };
+    const { result, rerender } = renderHook(
+      ({ value }) => useDeepCompareMemoize(value),
+      {
+        initialProps: { value: initial },
+      },
+    );
+    rerender({ value: { nested: [1, 2] } });
+    expect(result.current).toBe(initial);
+    const changed = { nested: [1, 3] };
+    rerender({ value: changed });
+    expect(result.current).toBe(changed);
+    rerender({ value: { nested: [1, 3] } });
+    expect(result.current).toBe(changed);
+  });
+
+  it("handles inline callbacks without causing additional renders", () => {
+    let renders = 0;
+    const { result, rerender } = renderHook(
+      ({ value }) => {
+        renders++;
+        return useDeepCompareMemoize({ callback: () => value });
+      },
+      { initialProps: { value: "initial" } },
+    );
+
+    const initialRenders = renders;
+    rerender({ value: "changed" });
+    expect(result.current.callback()).toBe("changed");
+    expect(renders).toBe(initialRenders + 1);
+  });
+});