@marimo-team/frontend 0.25.3-dev129 → 0.25.3-dev134

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

@@ -1,7 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { NotebookEvents } from "@marimo-team/client";
-import { afterEach, describe, expect, it, vi } from "vitest";
+import { createStore } from "jotai/vanilla";
+import { afterEach, describe, expect, it, onTestFinished, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { ConnectionState } from "../../host/connection";
+import { createNotebookReplica } from "../../host/replica";
 import { createEditorBackend } from "../backend";
 
 const wasm = vi.hoisted(() => ({ enabled: false, consumer: vi.fn() }));
@@ -16,16 +20,34 @@ afterEach(() => {
   vi.unstubAllGlobals();
 });
 
-describe("createEditorBackend", () => {
-  const options = {
-    static: false,
-    transportType: "websocket" as const,
-    url: () => "ws://localhost/ws",
-    headers: () => ({}),
-    notebookUrl: (id: string) =>
-      new URL(`http://localhost/api/marimo/v1/notebooks/${id}`),
-  };
+const options = {
+  static: false,
+  transportType: "websocket" as const,
+  url: () => "ws://localhost/ws",
+  headers: () => ({}),
+  notebookUrl: (id: string) =>
+    new URL(`http://localhost/api/marimo/v1/notebooks/${id}`),
+  waitUntilReady: async () => undefined,
+  checkHealth: async () => true,
+};
+
+/** Record each socket the transport opens. They never connect. */
+function sockets() {
+  const opened: string[] = [];
+  vi.stubGlobal(
+    "WebSocket",
+    class extends EventTarget {
+      public constructor(url: URL) {
+        super();
+        opened.push(String(url));
+      }
+      public close() {}
+    },
+  );
+  return opened;
+}
 
+describe("createEditorBackend", () => {
   it.each([
     ["a static notebook", { ...options, static: true }, "idle"],
     ["WASM", options, "wasm"],
@@ -34,18 +56,8 @@ describe("createEditorBackend", () => {
   ] as const)("follows %s", async (_name, chosen, expected) => {
     wasm.enabled = expected === "wasm";
     const fetch = vi.fn(() => new Promise<Response>(() => undefined));
-    const socket = vi.fn();
+    const opened = sockets();
     vi.stubGlobal("fetch", fetch);
-    vi.stubGlobal(
-      "WebSocket",
-      class extends EventTarget {
-        public constructor(url: URL) {
-          super();
-          socket(url);
-        }
-        public close() {}
-      },
-    );
     const events = new NotebookEvents(createEditorBackend(chosen).transport);
     await vi.waitFor(() =>
       expect(
@@ -54,7 +66,7 @@ describe("createEditorBackend", () => {
             events.status === "connected" && !wasm.consumer.mock.calls.length,
           wasm: wasm.consumer.mock.calls.length === 1,
           fetch: fetch.mock.calls.length === 1,
-          socket: socket.mock.calls.length === 1,
+          socket: opened.length === 1,
         }[expected],
       ).toBe(true),
     );
@@ -109,3 +121,80 @@ describe("createEditorBackend", () => {
     expect(fetch).not.toHaveBeenCalled();
   });
 });
+
+describe("the session backend's startup", () => {
+  it.each([true, false])(
+    "rechecks health after failed startup (healthy: %s)",
+    async (healthy) => {
+      const opened = sockets();
+      const waitUntilReady = vi
+        .fn()
+        .mockRejectedValue(new Error("Startup gave up"));
+      const checkHealth = vi.fn().mockResolvedValue(healthy);
+      const store = createStore();
+      const client = createNotebookReplica(store);
+      onTestFinished(
+        client.connect(
+          createEditorBackend({ ...options, waitUntilReady, checkHealth }),
+          { onNotification: vi.fn() },
+        ),
+      );
+      await vi.waitFor(() =>
+        expect(store.get(client.connectionAtom).state).toBe(
+          ConnectionState.CLOSED,
+        ),
+      );
+
+      await client.reconnect();
+      await vi.waitFor(() => expect(opened).toHaveLength(healthy ? 1 : 0));
+      expect({
+        healthChecks: checkHealth.mock.calls.length,
+        startupWaits: waitUntilReady.mock.calls.length,
+        state: store.get(client.connectionAtom).state,
+      }).toEqual({
+        healthChecks: 1,
+        startupWaits: 1,
+        state: healthy ? ConnectionState.CONNECTING : ConnectionState.CLOSED,
+      });
+    },
+  );
+
+  it("shares one startup wait across connections", async () => {
+    const opened = sockets();
+    const ready = new Deferred<void>();
+    const waitUntilReady = vi.fn(() => ready.promise);
+    const backend = createEditorBackend({ ...options, waitUntilReady });
+    const client = createNotebookReplica(createStore());
+    // Strict Mode connects, disconnects, and connects again.
+    client.connect(backend, { onNotification: vi.fn() })();
+    onTestFinished(client.connect(backend, { onNotification: vi.fn() }));
+    await Promise.resolve();
+    expect(opened).toEqual([]);
+
+    ready.resolve();
+    await vi.waitFor(() => expect(opened).toHaveLength(1));
+    expect(waitUntilReady).toHaveBeenCalledOnce();
+  });
+
+  it("opens nothing for a connection that ended during startup", async () => {
+    const opened = sockets();
+    const ready = new Deferred<void>();
+    const store = createStore();
+    const client = createNotebookReplica(store);
+    const disconnect = client.connect(
+      createEditorBackend({ ...options, waitUntilReady: () => ready.promise }),
+      { onNotification: vi.fn() },
+    );
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.CONNECTING,
+    });
+    disconnect();
+    ready.resolve();
+    await ready.promise;
+    await Promise.resolve();
+    expect(opened).toEqual([]);
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.NOT_STARTED,
+    });
+  });
+});

package/src/core/websocket/__tests__/useNotebookReplica.hook.test.tsx

@@ -0,0 +1,141 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import type { PropsWithChildren } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { NotebookTransport } from "@marimo-team/client";
+import { ConnectionState } from "../types";
+import { createEditorBackend } from "../backend";
+import { useNotebookReplica } from "../useNotebookReplica";
+import { fakeBackend, fixtureMessages } from "../../host/__tests__/fixtures";
+import { agentAttachmentsAtom } from "../../attachments/state";
+import { notebookReplicaAtom } from "../state";
+
+vi.mock("../backend", () => ({ createEditorBackend: vi.fn() }));
+
+/** A host-framed native notification. */
+const reconnected = {
+  event: "marimo",
+  data: { marimo: { op: "reconnected" } },
+};
+
+/** A transport whose connections the test drives. */
+function makeTransport() {
+  const opened: {
+    signal: AbortSignal;
+    open(): void;
+    receive(message: unknown): void;
+  }[] = [];
+  const transport: NotebookTransport = ({ signal }, handlers) =>
+    new Promise((resolve) => {
+      opened.push({ signal, open: handlers.open, receive: handlers.message });
+      signal.addEventListener("abort", () => resolve(), { once: true });
+    });
+  return {
+    transport,
+    opened,
+    /** Open the latest connection and deliver a host-framed notification. */
+    notify: () => opened.at(-1)?.receive(reconnected),
+  };
+}
+
+function setup(strict = false) {
+  const store = createStore();
+  const options = {
+    static: false,
+    transportType: "websocket" as const,
+    url: () => "ws://localhost/ws",
+    headers: () => ({}),
+    notebookUrl: (id: string) => new URL(`http://localhost/notebooks/${id}`),
+    waitUntilReady: async () => undefined,
+    checkHealth: async () => true,
+    onNotification: vi.fn(),
+  };
+  const wrapper = ({ children }: PropsWithChildren) => (
+    <Provider store={store}>{children}</Provider>
+  );
+  const hook = renderHook((props) => useNotebookReplica(props), {
+    initialProps: options,
+    reactStrictMode: strict,
+    wrapper,
+  });
+  return { ...hook, options, store };
+}
+
+beforeEach(() => {
+  vi.mocked(createEditorBackend).mockReset();
+});
+
+describe("useNotebookReplica", () => {
+  it("exposes host presence to the UI without waiting for kernel-ready", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { result, store, unmount } = setup();
+    await act(async () => Promise.resolve());
+    expect(store.get(notebookReplicaAtom)).toBe(result.current);
+    const [connection] = fake.opened;
+    act(() => {
+      connection.open();
+      fixtureMessages("notebook-events.sse").forEach(connection.receive);
+    });
+    expect(store.get(agentAttachmentsAtom)).toEqual([
+      {
+        id: "a2",
+        notebook_id: "n1",
+        kind: "agent",
+        name: "Notebook CLI",
+        since: "2026-10-01T09:20:03Z",
+      },
+    ]);
+    unmount();
+    expect(store.get(notebookReplicaAtom)).toBeNull();
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+
+  it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { rerender, options, unmount } = setup();
+    await act(async () => Promise.resolve());
+    const onNotification = vi.fn();
+    rerender({ ...options, onNotification, url: () => "ws://redirected/ws" });
+    act(() => fake.notify());
+    expect(options.onNotification).not.toHaveBeenCalled();
+    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
+      data: { op: "reconnected" },
+    });
+    const [{ url }] = vi.mocked(createEditorBackend).mock.calls[0];
+    expect(url()).toBe("ws://redirected/ws");
+    expect(fake.opened).toHaveLength(1);
+    unmount();
+  });
+
+  it("follows only the active connection under Strict Mode", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { result, store, options, unmount } = setup(true);
+    await act(async () => Promise.resolve());
+    const live = fake.opened.filter((attempt) => !attempt.signal.aborted);
+    expect(live).toHaveLength(1);
+    const [active] = live;
+    act(() => active.open());
+    expect(store.get(result.current.connectionAtom)).toEqual({
+      state: ConnectionState.OPEN,
+    });
+    act(() => fake.opened.forEach((attempt) => attempt.receive(reconnected)));
+    expect(options.onNotification).toHaveBeenCalledOnce();
+    unmount();
+    expect(active.signal.aborted).toBe(true);
+    const attempts = fake.opened.length;
+    await result.current.reconnect();
+    expect(fake.opened).toHaveLength(attempts);
+  });
+});

package/src/core/websocket/__tests__/useNotebookSession.hook.test.tsx

@@ -52,6 +52,7 @@ import { RuntimeManager } from "../../runtime/runtime";
 import { WebSocketClosedReason, ConnectionState } from "../types";
 import { classifyCloseEvent } from "../connection-state";
 import { createEditorBackend } from "../backend";
+import { fakeBackend } from "../../host/__tests__/fixtures";
 import { useNotebookSession } from "../useNotebookSession";
 import { type NotebookTransport, StreamStopped } from "@marimo-team/client";
 
@@ -85,10 +86,9 @@ type FakeTransport = ReturnType<typeof makeTransport>;
 
 /** Use a fake transport for the next session the hook follows. */
 function withTransport(fake: FakeTransport = makeTransport()) {
-  vi.mocked(createEditorBackend).mockReturnValue({
-    transport: fake.transport,
-    notebook: () => ({ edit: vi.fn() }),
-  });
+  vi.mocked(createEditorBackend).mockImplementation(({ checkHealth }) =>
+    fakeBackend({ transport: fake.transport, checkHealth }),
+  );
   return fake;
 }
 

package/src/core/websocket/backend.ts

@@ -1,8 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { NotebookClient, type StreamUrl } from "@marimo-team/client";
+/** The backend the editor follows, chosen once for its mode. */
+
+import { NotebookClient, StreamStopped } from "@marimo-team/client";
 import type { NotebookId } from "@marimo-team/client/protocol";
-import type { NotebookBackend } from "../host/backend";
+import type { NotebookBackend } from "../host";
 import { PyodideBridge } from "../wasm/bridge";
 import { isWasm } from "../wasm/utils";
 import {
@@ -18,11 +20,45 @@ export interface EditorBackendOptions {
   static: boolean;
   transportType: "websocket" | "sse";
   /** The session route, read before each connection. */
-  url: StreamUrl;
+  url: () => string;
   /** The runtime's headers, read before each request. */
   headers: () => Record<string, string>;
   /** The host URL of a notebook the stream reported. */
   notebookUrl: (notebookId: NotebookId) => URL;
+  /** Wait until the host can accept the first connection. */
+  waitUntilReady: () => Promise<void>;
+  /** Whether the host is reachable before an explicit reconnect. */
+  checkHealth: () => Promise<boolean>;
+}
+
+/**
+ * Read the session route once the host has started.
+ *
+ * Every connection waits for the same startup. Once it settles, even by
+ * giving up, connections read the route directly. An explicit reconnect
+ * checks the host's health instead.
+ */
+function afterStartup({
+  url,
+  waitUntilReady,
+}: Pick<EditorBackendOptions, "url" | "waitUntilReady">) {
+  let startup: Promise<void> | undefined;
+  let settled = false;
+  return async () => {
+    if (!settled) {
+      startup ??= waitUntilReady().finally(() => {
+        settled = true;
+      });
+      try {
+        await startup;
+      } catch (error) {
+        throw new StreamStopped("Healthy connection never made", {
+          cause: error,
+        });
+      }
+    }
+    return url();
+  };
 }
 
 /** Commands for a notebook the editor cannot change through its backend. */
@@ -55,14 +91,16 @@ export function createEditorBackend(
       notebook: uneditable,
     };
   }
+  const url = afterStartup(options);
   return {
     transport:
       options.transportType === "sse"
-        ? sessionSse(options.url, options.headers)
-        : websocket(options.url),
+        ? sessionSse(url, options.headers)
+        : websocket(url),
     notebook: (id) =>
       new NotebookClient(options.notebookUrl(id), {
         fetch: withHeaders(options.headers),
       }),
+    checkHealth: options.checkHealth,
   };
 }

package/src/core/websocket/connection-state.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { CloseDecision } from "../host/connection";
+import type { CloseDecision } from "../host";
 import { Logger } from "@/utils/Logger";
 import { ConnectionState, WebSocketClosedReason } from "./types";
 

package/src/core/websocket/state.ts

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { atom } from "jotai";
-import type { NotebookClient } from "../host/notebook-client";
+import type { NotebookReplica } from "../host";
 
-/** The client owned by the mounted notebook session. */
-export const notebookClientAtom = atom<NotebookClient | null>(null);
+/** The replica owned by the mounted notebook session. */
+export const notebookReplicaAtom = atom<NotebookReplica | null>(null);

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

@@ -7,8 +7,8 @@ import {
 } from "@marimo-team/client";
 import { createStore } from "jotai/vanilla";
 import { afterEach, describe, expect, it, onTestFinished, vi } from "vitest";
-import { createNotebookClient } from "../../../host/notebook-client";
-import { fixtureMessages } from "../../../host/__tests__/fixtures";
+import { createNotebookReplica } from "../../../host/replica";
+import { fakeBackend, fixtureMessages } from "../../../host/__tests__/fixtures";
 import { ConnectionState } from "../../../host/connection";
 import { Deferred } from "@/utils/Deferred";
 import { pyodide, sessionSse, websocket } from "../notebook";
@@ -76,7 +76,7 @@ function network(transportType: "websocket" | "sse", url: StreamUrl) {
     transportType === "websocket"
       ? websocket(url)
       : sessionSse(url, () => ({}));
-  const backend = { transport, notebook: () => ({ edit: vi.fn() }) };
+  const backend = fakeBackend({ transport });
   let index = 0;
   return {
     backend,
@@ -296,15 +296,10 @@ describe.each(["websocket", "sse"] as const)(
         const url = vi.fn(() => target);
         const connection = network(transportType, url);
         const store = createStore();
-        const { client, connect } = createNotebookClient(store);
-        const lifetime = new AbortController();
-        onTestFinished(() => lifetime.abort());
-        connect({
-          backend: connection.backend,
-          signal: lifetime.signal,
-          checkHealth: async () => true,
-          onNotification: vi.fn(),
-        });
+        const client = createNotebookReplica(store);
+        onTestFinished(
+          client.connect(connection.backend, { onNotification: vi.fn() }),
+        );
         const first = await connection.opened();
         const snapshot = fixtureMessages("notebook-events.sse");
         snapshot.forEach(first.receive);
@@ -363,15 +358,10 @@ describe.each(["websocket", "sse"] as const)(
       let target = `${base}?session_id=s1&file=a.py&access_token=a`;
       const connection = network(transportType, () => target);
       const store = createStore();
-      const { client, connect } = createNotebookClient(store);
-      const lifetime = new AbortController();
-      onTestFinished(() => lifetime.abort());
-      connect({
-        backend: connection.backend,
-        signal: lifetime.signal,
-        checkHealth: async () => true,
-        onNotification: vi.fn(),
-      });
+      const client = createNotebookReplica(store);
+      onTestFinished(
+        client.connect(connection.backend, { onNotification: vi.fn() }),
+      );
       const first = await connection.opened();
       fixtureMessages("notebook-events.sse").forEach(first.receive);
       await vi.waitFor(() =>

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

@@ -12,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,
@@ -408,7 +408,7 @@ export function useNotebookSession(opts: {
       variant: "danger",
     });
   };
-  const notebook = useNotebookClient({
+  const replica = useNotebookReplica({
     static: isStaticNotebook(),
     transportType,
     url: () =>
@@ -438,7 +438,7 @@ 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;
       }
@@ -469,16 +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,
     setKernelStartupError,
     updateStartupProgress,
   ]);
 
-  return { notebook, connection, reconnect: notebook.reconnect };
+  return { replica, connection, reconnect: replica.reconnect };
 }

package/src/plugins/impl/chat/chat-ui.tsx

@@ -495,7 +495,7 @@ export const Chatbot: React.FC<Props> = (props) => {
                     : undefined,
                 })}
               </div>
-              <div className="flex justify-end text-xs gap-2 invisible group-hover:visible">
+              <div className="flex justify-end text-xs gap-2 opacity-0 group-hover:opacity-100 focus-within:opacity-100">
                 <CopyClipboardIcon
                   value={textContent}
                   className="h-3 w-3"
@@ -503,6 +503,7 @@ export const Chatbot: React.FC<Props> = (props) => {
                 />
                 <button
                   type="button"
+                  aria-label="Delete message"
                   onClick={() => handleDelete(message.id)}
                   className="text-xs text-(--slate-9) hover:text-(--slate-11)"
                 >

package/src/plugins/impl/data-explorer/components/column-summary.tsx

@@ -60,7 +60,9 @@ export const ColumnSummary: React.FC<Props> = ({ schema }) => {
             field: field,
           });
           return (
-            <span
+            <button
+              type="button"
+              aria-pressed={selectedField === field}
               key={field}
               className={cn(
                 "hover:bg-muted self-start px-2 py-2 rounded flex flex-row gap-1 items-center cursor-pointer lg:justify-center text-sm truncate shrink-0 overflow-hidden",
@@ -81,7 +83,7 @@ export const ColumnSummary: React.FC<Props> = ({ schema }) => {
                   ({cardinality})
                 </span>
               )}
-            </span>
+            </button>
           );
         })}
         {hasMore && (