@marimo-team/frontend 0.25.3-dev12 → 0.25.3-dev121

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,
@@ -11,69 +11,78 @@ import type {
 import {
   createNotebookClient,
   type NotebookClient,
-  type NotebookConnection,
 } 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 & {
-  sessionId: string;
-  messageFormat: NotebookConnection["messageFormat"];
+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;
-  onMessageError: NotebookConnection["onMessageError"];
 };
 
 export function useNotebookClient(options: Options): NotebookClient {
   const store = useStore();
-  const { sessionId, static: isStatic, transportType, messageFormat } = options;
+  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);
-  const onMessageError = useEvent(options.onMessageError);
 
   useEffect(() => {
     const lifetime = new AbortController();
+    let needsInitialReadiness = true;
+    store.set(notebookClientAtom, controller.client);
     controller.connect({
-      createTransport: (url) =>
-        createConnectionTransport({
-          url,
-          headers,
-          static: isStatic,
-          transportType,
-        }),
-      url,
-      messageFormat,
+      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,
-      onMessageError,
       // The host validates the envelope, not the native notification contents.
       onNotification: (notification) =>
         onNotification({ data: notification as NotificationMessage }),
-      decodeJson: jsonParseWithSpecialChar,
-      classifyClose: classifyCloseEvent,
     });
-    return () => lifetime.abort();
+    return () => {
+      lifetime.abort();
+      if (store.get(notebookClientAtom) === controller.client) {
+        store.set(notebookClientAtom, null);
+      }
+    };
   }, [
     controller,
-    sessionId,
+    store,
     isStatic,
     transportType,
-    messageFormat,
     url,
     headers,
     waitUntilReady,
     checkHealth,
     onNotification,
-    onMessageError,
   ]);
 
   return controller.client;

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";
@@ -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 {
@@ -136,11 +135,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 +145,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 +179,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.
@@ -417,11 +410,8 @@ export function useNotebookSession(opts: {
     });
   };
   const notebook = useNotebookClient({
-    sessionId,
     static: isStaticNotebook(),
     transportType,
-    messageFormat:
-      notebookId && transportType === "websocket" ? "host" : "session",
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
@@ -434,7 +424,6 @@ export function useNotebookSession(opts: {
       }
     },
     checkHealth: () => runtimeManager.reconcileFromHealth(),
-    onMessageError,
     onNotification: (notification) => {
       try {
         handleNotification(notification);
@@ -452,7 +441,6 @@ export function useNotebookSession(opts: {
       if (status.state === ConnectionState.NOT_STARTED) {
         return;
       }
-      setAttachments([]);
       if (status.state === ConnectionState.OPEN) {
         updateStartupProgress(null);
         return; // kernel-ready, not the socket, establishes session readiness.
@@ -487,7 +475,6 @@ export function useNotebookSession(opts: {
     notebook,
     store,
     setConnection,
-    setAttachments,
     setKernelStartupError,
     updateStartupProgress,
   ]);

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);
+  });
+});